aymous/lfm2-8b-a1b-short-conv

LFM2.5-8B-A1B gated short-conv

Fused gated short-conv decode step for LiquidAI/LFM2.5-8B-A1B (Lfm2MoeShortConv, the LFM2 token-mixer on 18 of 24 layers, L_cache=3). Collapses the eager decode chain — chunk → B·x → 3-tap depthwise conv → C·conv_out → conv-state roll — into a single launch per layer (the fallback path when causal_conv1d is absent). The in_proj/out_proj GEMVs stay cuBLAS, outside the op.

Signature: lfm2_short_conv(BCx, conv_state, weight) -> (y, new_conv_state), where BCx = in_proj(x) laid out [B(H) | C(H) | x(H)] and conv_state is the [H, 3] cache; returns the gated output and the rolled state.

Part of a 3-kernel pack (with lfm2-8b-a1b-moe-experts and lfm2-8b-a1b-rms-norm) that makes the whole model 1.90× faster end-to-end on GB10.

Usage

from kernels import get_kernel
k = get_kernel("aymous/lfm2-8b-a1b-short-conv", version=1, trust_remote_code=True)
y, new_state = k.lfm2_short_conv(BCx, conv_state, conv_weight)   # single-token decode, bf16

aarch64, torch 2.11/2.12 × CUDA 12.6/12.8/13.0/13.2, archs 7.5–9.0 + 10.0/11.0/12.0/12.1+PTX. Benchmarked on GB10 (sm_121).