| --- |
| library_name: kernels |
| license: apache-2.0 |
| --- |
| |
| # decode-ops |
|
|
| Fused decode-step glue for LLM inference on CPUs, aarch64 first, loadable |
| through `kernels`: RMSNorm with a fused residual add, in-place RoPE, |
| SiLU-gate, greedy argmax, and deterministic temperature / top-k / top-p |
| sampling. The reference baselines are the equivalent torch op sequences and |
| the HuggingFace warper order, reproduced exactly. Companions: |
| [bitnet-cpu](https://huggingface.co/kernels/phanerozoic/bitnet-cpu) and |
| [quant-matmul](https://huggingface.co/kernels/phanerozoic/quant-matmul) |
| (linears), |
| [cpu-attn](https://huggingface.co/kernels/phanerozoic/cpu-attn) (int8 KV |
| attention). |
|
|
| Once the matmuls and attention of a CPU decode step are fast, what remains |
| is the glue: norms, rotary, gating, and sampling, each a chain of small |
| torch dispatches with intermediate tensors, and on a single-board computer |
| that chain is a real share of the step. These ops replace each chain with |
| one call, and the sampler is deterministic by construction: an explicit |
| uniform draw and a fixed survivor order mean a given generator always yields |
| the same token. |
|
|
|  |
|
|
| *Measured on the boards: RMSNorm at 4096 in 15 us against torch's 38 (2.4x), |
| argmax over a 128,256 vocabulary in 109 us against 597 (5.5x), and full |
| top-k/top-p sampling in 1.6 ms where the torch sequence has no single |
| equivalent; the same kernel on a Pi 4's Cortex-A72 runs 76, 225, and |
| 7,862 us.* |
|
|
| ## Usage |
|
|
| ```python |
| import torch |
| from kernels import get_kernel |
| |
| do = get_kernel("phanerozoic/decode-ops", version=1, trust_remote_code=True) |
| |
| h = do.add_rmsnorm(residual, x, norm_weight) # residual += x, then norm |
| do.rope(q, k, pos) # in place |
| y = do.silu_mul(gate, up) |
| tok = do.sample(logits, temperature=0.8, top_k=50, top_p=0.95, |
| generator=gen) # deterministic given gen |
| tok = do.argmax(logits) # greedy |
| ``` |
|
|
| `version` selects the release branch; `trust_remote_code` is required by |
| `kernels` for publishers without the trusted-publisher mark. |
|
|
| ## API |
|
|
| | Symbol | Purpose | |
| |---|---| |
| | `rmsnorm(x, weight, eps)` | `x * w / sqrt(mean(x^2) + eps)` | |
| | `add_rmsnorm(residual, x, weight, eps)` | `residual += x` in place, then norm | |
| | `rope(q, k, pos, theta=10000, interleaved=False)` | rotate one position in place | |
| | `silu_mul(gate, up)` | `silu(gate) * up` | |
| | `argmax(logits)` | greedy token, first index on ties | |
| | `sample(logits, temperature, top_k, top_p, generator)` | deterministic sampling; `temperature <= 0` is greedy | |
|
|
| ## Method |
|
|
| RMSNorm accumulates the square sum in f32; RoPE precomputes its angle pairs |
| in f64 per call and rotates q and k in place (both half-split and |
| interleaved layouts); SiLU and softmax use a NEON polynomial exp with |
| `exp(0) == 1` exactly. Sampling follows the HF warper order (temperature, |
| then top-k, then top-p over the renormalized survivors) with an explicit |
| uniform draw: survivors pop from a heap in descending logit order, ties to |
| the lower index, and the token is the inverse CDF at the draw. RMSNorm, the |
| residual add, and SiLU-gate run on NEON (aarch64) or AVX2 + FMA (x86-64), |
| selected once at runtime; the sampling and argmax exit paths are scalar. |
| `DO_CPU_ISA` (`scalar`, `avx2`, `neon`) demotes for A/B runs. |
|
|
| ## Measured |
|
|
| Raspberry Pi 5 (4x Cortex-A76 2.4 GHz) and Pi 4 Model B (4x Cortex-A72 |
| 1.8 GHz), torch 2.13 CPU, against the equivalent torch op sequences: |
|
|
| | op | Pi 5 | torch (Pi 5) | Pi 4 | torch (Pi 4) | |
| |---|---|---|---|---| |
| | rmsnorm (1 x 4096) | 15 us | 38 us | 76 us | 576 us | |
| | silu_mul (1 x 6912) | 24 us | 34 us | 162 us | 156 us | |
| | argmax (V=128256) | 109 us | 597 us | 225 us | 1,435 us | |
| | sample (V=128256, k=50, p=0.95) | 1,596 us | no equivalent | 7,862 us | no equivalent | |
| |
| Further ops on the same boards: `add_rmsnorm` (1 x 4096) runs 19 us on the |
| Pi 5 and 75 us on the Pi 4 against torch's 46 and 171; `rope` over 32 + 8 |
| heads at dim 128 runs 5.4 and 15 us. |
|
|
| ## Correctness |
|
|
| RMSNorm and SiLU-gate match the torch expressions elementwise; argmax is |
| exact against `torch.argmax` including tie behavior. Sampling reproduces the |
| HF warper order and is deterministic given the generator state: a fixed |
| generator yields a fixed token, verified across both boards. |
|
|
| ## Requirements and limits |
|
|
| - f32 activations (bf16 inputs are converted at the wrapper). |
| - x86-64 runs the scalar path for the sampling and argmax exits: correct, |
| not tuned. |
| - SiLU-gate on the Pi 4 is at parity with torch; the wins are in the norm, |
| argmax, and sampling paths. |
|
|
| ## References |
|
|
| Zhang and Sennrich, "Root Mean Square Layer Normalization" (2019); Su et |
| al., "RoFormer" (2021); the HuggingFace logits-warper order as the sampling |
| reference. |
|
|
| ## License |
|
|
| Apache-2.0. |
|
|
