aymous/lfm2-8b-a1b-moe-experts

LFM2.5-8B-A1B MoE experts

Fused MoE routed experts for LiquidAI/LFM2.5-8B-A1B (lfm2_moe, 8.5B total / 1B active): 32 experts, top-4, plain SiLU SwiGLU, router-weight scale + scatter-add. The router (Lfm2MoeSparseMoeBlock: sigmoid + top-4, normalized by sum) stays outside the op; no shared expert. Two entry points on (hidden_states, top_k_index, top_k_weights): lfm2moe_moe_experts (grouped GEMM, prefill) and lfm2moe_moe_experts_decode (fused gather-GEMV for n_tokens ~1–4). Reference op: transformers Lfm2MoeExperts.

Part of a 3-kernel pack (with lfm2-8b-a1b-short-conv and lfm2-8b-a1b-rms-norm) that makes the whole model 1.90× faster end-to-end (28.5 → 54.3 tok/s decode on GB10); this MoE kernel is the dominant lever (1.76×).

Benchmark (NVIDIA GB10 / sm_121, bf16, blocked_autorange)

workload	op_eager	op_compile	custom	speedup
decode (1 tok)	0.863 ms	0.917 ms	0.396 ms	2.31× vs compile
prefill (2048)	8.636 ms	8.786 ms	6.304 ms	1.39× vs compile

Verified vs Lfm2MoeExperts (atol=rtol=3e-2; fp32 kernel accumulation vs bf16 reference).

Usage

from kernels import get_kernel
k = get_kernel("aymous/lfm2-8b-a1b-moe-experts", version=1, trust_remote_code=True)
out = k.lfm2moe_moe_experts_decode(x, gate_up_proj, down_proj, top_k_index, top_k_weights, top_k)

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.