Views
No views yet
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.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 |
Lfm2MoeExperts (atol=rtol=3e-2; fp32 kernel accumulation vs bf16 reference).1from kernels import get_kernel
2k = get_kernel("aymous/lfm2-8b-a1b-moe-experts", version=1, trust_remote_code=True)
3out = 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.