Kernels
kernel
File size: 1,742 Bytes
cf5d755
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
---
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](https://huggingface.co/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](https://huggingface.co/aymous/lfm2-8b-a1b-short-conv) and [lfm2-8b-a1b-rms-norm](https://huggingface.co/aymous/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
```python
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.