[E2E Model Support] Add TileGym kernel integration for LFM2 MoE

modeling/transformers/scripts/benchmark_hf_model.sh

@@ -22,7 +22,7 @@ usage() {
     cat <<EOF
 Usage: $0 --model-key KEY [options]
 
-Model keys: llama, deepseek, qwen, qwen3_5, gemma3, gpt_oss, mistral, phi3, olmo3, olmoe
+Model keys: llama, deepseek, qwen, qwen3_5, gemma3, gpt_oss, mistral, phi3, olmo3, olmoe, lfm2_moe
 
 Options:
   --model-id ID           Override Hugging Face model id or local model path
@@ -145,6 +145,14 @@ case "${MODEL_KEY}" in
         DEFAULT_SUMMARY_FILE="${LOG_DIR}/olmoe_benchmark_summary.txt"
         TITLE="OLMoE-1B-7B-0924"
         ;;
+    lfm2_moe)
+        DEFAULT_MODEL_ID="LiquidAI/LFM2-8B-A1B"
+        DEFAULT_INPUT_FILE="${PROJECT_DIR}/sample_inputs/input_prompt_small.txt"
+        DEFAULT_OUTPUT_LENGTH=50
+        DEFAULT_BATCH_SIZE=1
+        DEFAULT_SUMMARY_FILE="${LOG_DIR}/lfm2_moe_benchmark_summary.txt"
+        TITLE="LFM2-8B-A1B"
+        ;;
     *)
         echo "Unknown --model-key: ${MODEL_KEY}"
         usage

modeling/transformers/src/tilegym_hf_bench/tilegym_patch.py

@@ -5,6 +5,7 @@
 from tilegym.transformers import apply_tilegym_kernel_to_deepseek_v2
 from tilegym.transformers import apply_tilegym_kernel_to_gemma3
 from tilegym.transformers import apply_tilegym_kernel_to_gpt_oss
+from tilegym.transformers import apply_tilegym_kernel_to_lfm2_moe
 from tilegym.transformers import apply_tilegym_kernel_to_llama
 from tilegym.transformers import apply_tilegym_kernel_to_mistral
 from tilegym.transformers import apply_tilegym_kernel_to_olmo3
@@ -42,5 +43,9 @@ def apply_tilegym_patch(model_id, use_attn=False, use_cutile=False):
         apply_tilegym_kernel_to_olmoe(rope=True, rms_norm=True, attn=use_attn, moe=True, use_cutile=use_cutile)
     elif "olmo-3" in model_name or "olmo3" in model_name:
         apply_tilegym_kernel_to_olmo3(rope=True, rms_norm=True, swiglu=True, attn=use_attn, use_cutile=use_cutile)
+    elif "lfm2" in model_name and ("moe" in model_name or "a1b" in model_name):
+        # LFM2-MoE checkpoints (e.g. LiquidAI/LFM2-8B-A1B). The dense LFM2 family is
+        # not supported, so only match the MoE variants (…moe / …A1B).
+        apply_tilegym_kernel_to_lfm2_moe(rope=True, rms_norm=True, attn=use_attn, moe=True, use_cutile=use_cutile)
     else:
         print(f"Warning: Model {model_id} is not supported in tilegym patch. No optimizations will be applied.")

src/tilegym/transformers/__init__.py

@@ -5,6 +5,7 @@
 from tilegym.transformers.monkey_patch import apply_tilegym_kernel_to_deepseek_v2
 from tilegym.transformers.monkey_patch import apply_tilegym_kernel_to_gemma3
 from tilegym.transformers.monkey_patch import apply_tilegym_kernel_to_gpt_oss
+from tilegym.transformers.monkey_patch import apply_tilegym_kernel_to_lfm2_moe
 from tilegym.transformers.monkey_patch import apply_tilegym_kernel_to_llama
 from tilegym.transformers.monkey_patch import apply_tilegym_kernel_to_mistral
 from tilegym.transformers.monkey_patch import apply_tilegym_kernel_to_olmo3

src/tilegym/transformers/lfm2_moe/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: MIT

src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_prefill.json

@@ -0,0 +1,58 @@
+{
+  "axes": {
+    "B": {
+      "description": "Supported batch size.",
+      "type": "const",
+      "value": 1
+    },
+    "D": {
+      "description": "Channel dimension.",
+      "type": "var"
+    },
+    "K": {
+      "description": "Convolution kernel size (conv_L_cache).",
+      "type": "const",
+      "value": 3
+    },
+    "T": {
+      "description": "Sequence length.",
+      "type": "var"
+    }
+  },
+  "description": "LFM2-MoE prefill-path depthwise causal conv1d (kernel_size=3, no activation). Input is unpadded; the entry point left-pads by K-1 internally.",
+  "inputs": {
+    "weight": {
+      "dtype": "bfloat16",
+      "shape": [
+        "D",
+        "K"
+      ]
+    },
+    "x": {
+      "dtype": "bfloat16",
+      "shape": [
+        "B",
+        "D",
+        "T"
+      ]
+    }
+  },
+  "name": "lfm2_moe_causal_conv1d_prefill",
+  "op_type": "causal_conv1d_prefill",
+  "outputs": {
+    "output": {
+      "dtype": "bfloat16",
+      "shape": [
+        "B",
+        "D",
+        "T"
+      ]
+    }
+  },
+  "reference": "# Source: https://github.com/huggingface/transformers/blob/5eddc12edfaf8cafde8c9bae4ccb12f8a139b4f9/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py#L385-L405\nimport torch\nimport torch.nn.functional as F\n\ndef run(x, weight, bias=None, activation=None, seq_idx=None):\n    # LFM2 short-conv fuses no activation and does not use packed sequences,\n    # so `activation` and `seq_idx` are always None at the call site.\n    _, hidden_size, seq_len = x.shape\n    padding = weight.shape[-1] - 1\n    out = F.conv1d(\n        x.to(weight.dtype),\n        weight=weight.unsqueeze(1),\n        bias=bias,\n        padding=padding,\n        groups=hidden_size,\n    )[:, :, :seq_len]\n    return out.to(x.dtype)",
+  "tags": [
+    "model:lfm2_moe",
+    "stage:prefill",
+    "status:unverified"
+  ]
+}

src/tilegym/transformers/lfm2_moe/kernel_definitions/lfm2_moe_causal_conv1d_update.json

@@ -0,0 +1,67 @@
+{
+  "axes": {
+    "B": {
+      "description": "Supported batch size.",
+      "type": "const",
+      "value": 1
+    },
+    "D": {
+      "description": "Channel dimension.",
+      "type": "var"
+    },
+    "K": {
+      "description": "Convolution kernel size / cached window length (conv_L_cache).",
+      "type": "const",
+      "value": 3
+    },
+    "T": {
+      "description": "Sequence length (single-token decode path).",
+      "type": "const",
+      "value": 1
+    }
+  },
+  "description": "LFM2-MoE decode-path depthwise causal conv1d update (kernel_size=3, no activation). The conv_state input holds the full K-wide window and is rolled/updated in place.",
+  "inputs": {
+    "conv_state": {
+      "dtype": "bfloat16",
+      "shape": [
+        "B",
+        "D",
+        "K"
+      ]
+    },
+    "weight": {
+      "dtype": "bfloat16",
+      "shape": [
+        "D",
+        "K"
+      ]
+    },
+    "x": {
+      "dtype": "bfloat16",
+      "shape": [
+        "B",
+        "D",
+        "T"
+      ]
+    }
+  },
+  "name": "lfm2_moe_causal_conv1d_update",
+  "op_type": "causal_conv1d_update",
+  "outputs": {
+    "output": {
+      "dtype": "bfloat16",
+      "shape": [
+        "B",
+        "D",
+        "T"
+      ]
+    }
+  },
+  "reference": "# Source: https://github.com/huggingface/transformers/blob/5eddc12edfaf8cafde8c9bae4ccb12f8a139b4f9/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py#L365-L382\nimport torch\nimport torch.nn.functional as F\n\ndef run(x, conv_state, weight, bias=None, activation=None):\n    # `conv_state` holds the full K-wide window and is rolled in place.\n    # LFM2 short-conv fuses no activation, so `activation` is always None.\n    _, hidden_size, seq_len = x.shape\n    state_len = conv_state.shape[-1]\n    hidden_states_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype)\n    conv_state.copy_(hidden_states_new[:, :, -state_len:])\n    out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)\n    return out[:, :, -seq_len:].to(x.dtype)",
+  "tags": [
+    "model:lfm2_moe",
+    "stage:decode",
+    "status:unverified"
+  ]
+}

src/tilegym/transformers/lfm2_moe/kernel_solutions/lfm2_moe_causal_conv1d_prefill.json

@@ -0,0 +1,25 @@
+{
+  "author": "tilegym-agent",
+  "definition": "lfm2_moe_causal_conv1d_prefill",
+  "description": "cuTile prefill-path depthwise causal conv1d (kernel_size=3, no activation) for LFM2-MoE.",
+  "name": "lfm2_moe_causal_conv1d_prefill_cutile",
+  "sources": {
+    "path": [
+      "src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_prefill.py"
+    ]
+  },
+  "spec": {
+    "dependencies": [
+      "torch",
+      "cuda-tile"
+    ],
+    "destination_passing_style": false,
+    "entry_point": "src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_prefill.py::lfm2_causal_conv1d_fn_cutile",
+    "language": "cuda-tile",
+    "target_hardware": [
+      "SM100",
+      "SM103",
+      "SM120"
+    ]
+  }
+}

src/tilegym/transformers/lfm2_moe/kernel_solutions/lfm2_moe_causal_conv1d_update.json

@@ -0,0 +1,25 @@
+{
+  "author": "tilegym-agent",
+  "definition": "lfm2_moe_causal_conv1d_update",
+  "description": "cuTile decode-path depthwise causal conv1d update (kernel_size=3, no activation) for LFM2-MoE. Updates conv_state in place.",
+  "name": "lfm2_moe_causal_conv1d_update_cutile",
+  "sources": {
+    "path": [
+      "src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py"
+    ]
+  },
+  "spec": {
+    "dependencies": [
+      "torch",
+      "cuda-tile"
+    ],
+    "destination_passing_style": false,
+    "entry_point": "src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_update.py::lfm2_causal_conv1d_update_cutile",
+    "language": "cuda-tile",
+    "target_hardware": [
+      "SM100",
+      "SM103",
+      "SM120"
+    ]
+  }
+}

src/tilegym/transformers/lfm2_moe/kernels/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: MIT

src/tilegym/transformers/lfm2_moe/kernels/causal_conv1d_prefill.py

@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+
+"""LFM2-MoE depthwise causal conv1d prefill-path cuTile kernel.
+
+Replacement for the module-level ``causal_conv1d_fn`` called by
+``Lfm2MoeShortConv.forward``. Unlike the Qwen3.5 conv kernel this
+uses ``kernel_size = conv_L_cache = 3`` and applies **no** activation — LFM2's
+short conv is externally gated (``y = C * conv_out``), the SiLU is not fused in.
+
+The stock call site passes the *unpadded* input ``Bx`` of shape ``(B, C, L)``;
+the wrapper left-pads by ``K-1`` and the kernel reads the ``K``-wide causal
+window, matching ``nn.Conv1d(..., padding=L_cache-1)(Bx)[..., :L]``.
+"""
+
+import cuda.tile as ct
+import torch
+import torch.nn.functional as F
+
+ConstInt = ct.Constant[int]
+
+
+@ct.kernel
+def _causal_conv1d_prefill_kernel(
+    x,  # (D, T_padded)  left-padded by K-1
+    weight,  # (D, K=3)
+    output,  # (D, T)
+    BLOCK_T: ConstInt,
+):
+    bid_d = ct.bid(0)
+    bid_t = ct.bid(1)
+    t_start = bid_t * BLOCK_T
+    offs = ct.arange(BLOCK_T, dtype=ct.int32)
+    t_idx = t_start + offs
+
+    w0 = ct.astype(ct.gather(weight, (bid_d, 0), check_bounds=True), ct.float32)
+    w1 = ct.astype(ct.gather(weight, (bid_d, 1), check_bounds=True), ct.float32)
+    w2 = ct.astype(ct.gather(weight, (bid_d, 2), check_bounds=True), ct.float32)
+
+    # x is left-padded, so window for output position t reads x[t], x[t+1], x[t+2].
+    v0 = ct.astype(ct.gather(x, (bid_d, t_idx), padding_value=0.0, check_bounds=True), ct.float32)
+    v1 = ct.astype(ct.gather(x, (bid_d, t_idx + 1), padding_value=0.0, check_bounds=True), ct.float32)
+    v2 = ct.astype(ct.gather(x, (bid_d, t_idx + 2), padding_value=0.0, check_bounds=True), ct.float32)
+
+    result = v0 * w0 + v1 * w1 + v2 * w2
+
+    ct.scatter(output, (bid_d, t_idx), ct.astype(result, output.dtype), check_bounds=True)
+
+
+def lfm2_causal_conv1d_fn_cutile(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    activation=None,
+    seq_idx=None,
+) -> torch.Tensor:
+    """Depthwise causal conv1d for the prefill path (drop-in for ``causal_conv1d_fn``).
+
+    Args:
+        x: ``(B=1, D, L)`` unpadded input (``Bx`` in the LFM2 short-conv block).
+        weight: ``(D, K=3)`` depthwise conv weights (``conv.weight.view(D, K)``).
+        bias: optional ``(D,)`` bias (LFM2-8B-A1B uses ``conv_bias=False`` -> None).
+        activation: accepted for signature compatibility; must be ``None`` (LFM2
+            does not fuse an activation into the conv).
+        seq_idx: accepted for signature compatibility. Packed-sequence boundaries
+            are not supported by this fused path; only ``None`` is handled.
+
+    Returns:
+        ``(B=1, D, L)`` conv output.
+    """
+    assert activation is None, "LFM2 short-conv fuses no activation; activation must be None"
+    assert seq_idx is None, "lfm2_causal_conv1d_fn_cutile does not support packed sequences (seq_idx)"
+
+    B, D, L = x.shape
+    assert B == 1, "lfm2_causal_conv1d_fn_cutile only supports B=1"
+    K = weight.shape[1]
+    assert K == 3, f"expected kernel_size 3, got {K}"
+
+    x_2d = x.squeeze(0).contiguous()  # (D, L)
+    x_padded = F.pad(x_2d, (K - 1, 0))  # (D, L + K - 1), left pad only
+    w = weight.contiguous()
+    output = torch.empty(D, L, dtype=x.dtype, device=x.device)
+
+    # NOTE: the sequence length is deliberately *not* passed as a ct.Constant --
+    # it would become part of the JIT specialization key and force a fresh cuTile
+    # compile for every distinct prompt length (a long-prompt stall). Bounds are
+    # handled by `check_bounds` on the gathers/scatter, so one compiled kernel
+    # serves all lengths.
+    BLOCK_T = 256
+    grid = (D, (L + BLOCK_T - 1) // BLOCK_T)
+    ct.launch(
+        torch.cuda.current_stream(),
+        grid,
+        _causal_conv1d_prefill_kernel,
+        (x_padded, w, output, BLOCK_T),
+    )
+
+    out = output.unsqueeze(0)  # (1, D, L)
+    if bias is not None:
+        out = out + bias.view(1, -1, 1)
+    return out