Kernels:

---

phanerozoic

/

decode-ops

like 0

---
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.

![Per-op bars comparing torch, the kernel on a Pi 5, and the kernel on a Pi 4](https://huggingface.co/kernels/phanerozoic/decode-ops/resolve/main/media/hero.gif)

*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.