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 and quant-matmul (linears), 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

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

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:

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.