tags:
- kernel
library_name: kernels
license: apache-2.0
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.