decode-ops

Fused decode-step glue for LLM inference on CPUs, aarch64-first: RMSNorm with a fused residual add, in-place RoPE, SiLU-gate, greedy argmax, and deterministic temperature / top-k / top-p sampling. This is the non-matmul residue that owns a CPU decode step once the linears and attention are fast. Companions on the Hub: phanerozoic/bitnet-cpu and phanerozoic/quant-matmul (linears), phanerozoic/cpu-attn (int8 KV attention).

Each op replaces a chain of small torch CPU dispatches and their intermediate tensors with one call. 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, top-k, then top-p over the renormalized survivors) with an explicit uniform draw, so a fixed generator gives a fixed token: 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.

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

API

Performance

Raspberry Pi 5 (4x Cortex-A76 2.4 GHz) and Raspberry Pi 4 Model B (4x Cortex-A72 1.8 GHz), 64-bit Raspberry Pi OS, torch 2.13 CPU, against the equivalent torch op sequences:

Requirements

• f32 activations (bf16 inputs are converted at the wrapper).
• Sampling reproduces the HF warper order; results are deterministic given the generator state.
• x86-64 runs the scalar path: correct, not tuned.

Scope

The glue half of a CPU decode stack. The roadmap item is a fused decoder_layer_forward chaining norm, linears, RoPE, attention, and the FFN through the companion kernels in one call per layer.

License

Apache-2.0.