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

[iamanishx]

Description

Adds end-to-end TileGym support for LiquidAI/LFM2-8B-A1B. Closes #154.

LFM2-8B-A1B is a hybrid MoE model, 8.3B total and 1.5B active parameters. In this
checkpoint the token mixer is a short depthwise causal convolution in 18 of 24 layers
and GQA attention in the other 6, while the channel mixer is a dense SwiGLU MLP in the
first 2 layers and a 32-expert top-4 sparse MoE block in the remaining 22.

TileGym had no causal conv1d kernel for the convolution blocks, so this PR contributes
two new cuTile kernels alongside the model integration, following the layout of
src/tilegym/transformers/qwen3_5/kernels/. The MoE block follows
OlmoeSparseMoeBlockTileGym, per the guidance on the issue.

Integration. apply_tilegym_kernel_to_lfm2_moe in monkey_patch.py, registered as
lfm2_moe in MODEL_TYPE_TO_APPLY_TILEGYM_FN:

Component	Replacement
RoPE	get_apply_rope_func(model="llama")
RMSNorm	get_rms_norm_module(), operator, ffn, embedding and per-head Q/K norms
Attention	get_fmha_interface() on ALL_ATTENTION_FUNCTIONS["sdpa"]
MoE block	Lfm2MoeSparseMoeBlockTileGym over TileGym fused_moe
Dense MLP	Lfm2MoeMLPTileGym with fused silu_and_mul
Short conv	Two new cuTile depthwise causal conv1d kernels, gated on use_cutile=True

New files under src/tilegym/transformers/lfm2_moe/: kernels/causal_conv1d_prefill.py,
kernels/causal_conv1d_update.py, matching kernel_definitions/ and kernel_solutions/
pairs, and modeling_lfm2_moe.py. Also a lfm2_moe preset in
modeling/transformers/scripts/benchmark_hf_model.sh, dispatch in tilegym_patch.py,
tests/ops/test_causal_conv1d_lfm2.py, and the in-place conv_state declaration in
tests/kernel_inventory/runtime_inputs.yaml.

Weight layout and routing, verified against the HF source. gate_up_proj is
(E, 2I, H) with gate first then up, confirmed by HF's
linear(x, gate_up_proj[e]).chunk(2, dim=-1), and down_proj is (E, H, I). Both match
what fused_moe expects, so weights pass through with no merge and no reorder. Routing
uses sigmoid rather than softmax. When use_expert_bias is set the bias only selects the
top-k experts while the gathered weights stay unbiased. norm_topk_prob divides by
sum + 1e-6, then routed_scaling_factor scales. Submodule and buffer names are
identical to the stock block, so the HF state dict loads with strict=True, and
Lfm2MoeMLPTileGym keeps the w1, w3, w2 names LFM2 uses.

transformers compatibility. The LFM2 MoE API differs across releases and both are
handled. On 5.13 and later the conv entry points are module-level functions called from
Lfm2MoeShortConv.forward with 3D (B, D, 1) decode tensors, and the router is a
dedicated Lfm2MoeTopKRouter. On 5.10.x the call comes from cuda_kernels_forward with
2D Bx.squeeze(-1), and the router is inlined in Lfm2MoeSparseMoeBlock. The update
wrapper accepts both ranks and returns the matching rank; the MoE wrapper falls back to
nn.Linear, which exposes the same gate.weight state dict key.

Testing on RTX PRO 6000 Blackwell, SM120, the same compute capability as the RTX 5090
referenced in the issue, with transformers 5.10.2 and torch 2.9.1+cu130:

pytest tests/ops/test_causal_conv1d_lfm2.py -k test_op   17 passed
pytest tests/kernel_inventory -k lfm2                     2 passed

Kernel coverage: prefill against F.conv1d in float32 and bfloat16 over several channel
and sequence sizes, decode update against a torch reference including the in-place state
roll, the legacy 2D decode call, and prefill against token-by-token decode for
consistency. End-to-end generation was run unpatched, patched without conv kernels, and
patched with conv kernels; all three produce coherent output. Outputs are not
bit-identical across configurations, which is expected because sigmoid top-k routing
turns small numeric differences into different expert selections. Patch application was
confirmed on the loaded model: backend cutile, both conv entry points bound to the
TileGym wrappers, Lfm2MoeSparseMoeBlockTileGym, Lfm2MoeMLPTileGym, _TileRMSNorm,
sdpa, 18 of 24 layers on the conv path, and conv update call counts equal to conv
layers times decode tokens.

Performance, steady state after two warmup passes at the measured shape, five timed
runs, batch 1, bfloat16. Prefill, single forward pass:

Sequence length	Baseline	TileGym cuTile	Speedup
1024	40.8 ms	27.8 ms	1.47x
4096	82.4 ms	76.5 ms	1.08x
16384	318.4 ms	284.7 ms	1.12x

Decode, 210 token prompt, 64 generated tokens: baseline 62.30 tok/s, TileGym without conv
kernels 50.77 tok/s, TileGym with conv kernels 54.28 tok/s. Batch-1 decode is slower than
baseline for both TileGym configurations. That regime sits about 10x off the memory
roofline, roughly 16 ms per token against a 1.7 ms floor for 1.5B active parameters at
1792 GB/s, so it is bound by Python dispatch and per-call fixed costs rather than kernel
throughput; fused_moe pays token-to-expert sorting, alignment and workspace setup on
every call, which amortizes over thousands of tokens but not over one, while the stock
path only touches the four selected experts. The conv kernels recover 6.9 percent of
decode throughput, consistent with 18 of 24 layers using the conv path.

Notes. The conv kernels support batch size 1 and kernel_size 3, which is
conv_L_cache for this checkpoint; both are asserted. The attention patch registers on
ALL_ATTENTION_FUNCTIONS["sdpa"] only, matching most existing model patches, so loading
with attn_implementation="eager" bypasses FMHA; happy to also patch eager as gemma3
does. The conv patch requires use_cutile=True because the kernels are cuTile only.
Kernel definitions are tagged status:unverified; happy to flip them once the team's own
validation passes.

CI Configuration

config:
  build: true
  # valid options are "ops", "benchmark", and "sanity"
  test: ["ops", "sanity"]

Checklist

• Code formatted and imports sorted via repo specifications (./format.sh)
• Documentation updated (if needed)
• CI configuration reviewed

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hannahli-nv]

/ok to test b176ea4

[hannahli-nv]

Hi @iamanishx, thanks a lot for this PR! It's a really thorough end-to-end LFM2-MoE integration. We'll do a full review pass shortly.

One thing before we can merge: since this is your first contribution to TileGym, we'll need a signed CLA on file. Per first-time-contributors-cla-required, changes are MIT-licensed and require a signed CLA (LICENSES/CLA.md) emailed to TileGym@nvidia.com. Thanks for understanding!

[iamanishx]

One thing before we can merge: since this is your first contribution to TileGym, we'll need a signed CLA on file. Per first-time-contributors-cla-required, changes are MIT-licensed and require a signed CLA (LICENSES/CLA.md) emailed to TileGym@nvidia.com. Thanks for understanding!

sure

[iamanishx]

@hannahli-nv @lirundong hi, please have a look!!

[hannahli-nv]

@hannahli-nv @lirundong hi, please have a look!!

Apologies for the delayed response. We’ll update the review comments early next week. Thank you for your understanding regarding the inconvenience.