spec-verify

Speculative-decoding verification ops for CPUs, aarch64 first, loadable through kernels. The reference standard is the analytic law of rejection sampling, matched to the sampling floor. Companions: bitnet-cpu, quant-matmul, gguf-gemv (linears), cpu-attn, decode-ops.

On a small board the decode wall is memory: every token re-reads every weight, so one token costs the same bandwidth as many. Speculative decoding converts that into a win by having a draft model propose K tokens and the target verify them in a single [K+1, V] forward, which on the companion linear kernels is several times cheaper per token than K separate decodes. These ops are the verification step: an accepted prefix plus one correction token, distribution-preserving, at a cost far below the forwards they gate.

600 verification rounds at K = 4 measured on a Raspberry Pi with a deliberately imperfect draft: the accepted-length histogram gives a mean of 1.84 against the analytic law's 1.86, at 33 us greedy and 117 us sampled per round on a Pi 5 (126 and 657 us on a Pi 4), far below the model forwards they gate.

Usage

import torch
from kernels import get_kernel

sv = get_kernel("phanerozoic/spec-verify", version=1, trust_remote_code=True)

# Greedy: target_probs [K+1, V] (only the argmax is read), draft_tokens [K].
n_accepted, next_token = sv.verify_greedy(target_probs, draft_tokens)
# Keep draft_tokens[:n_accepted], append next_token.

# Or drive the whole draft/verify loop:
dec = sv.SpeculativeDecoder(draft_step, target_step, k=4, greedy=True)
out = dec.generate(prompt_ids, max_new=128)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

draft_step(prefix, k) -> (tokens, probs) and target_step(prefix, proposed) -> target_probs [k+1, V] are user callbacks, so any pair of models (including GGUF drafts and verifiers via the companion kernels) composes.

Method

Two verification modes:

• Greedy (temperature 0): accept draft token t_i while it equals argmax(target_i); at the first miss emit argmax(target_i) and stop; all K matched emit argmax(target_K). Deterministic, and the emitted stream is identical to plain greedy target decoding.
• Speculative sampling (Leviathan et al. / Chen et al.): accept t_i with probability min(1, p_i(t_i) / q_i(t_i)); on rejection resample from norm(max(0, p_i - q_i)); all accepted samples the bonus from p_K. Distribution-preserving when target and draft distributions are warped identically (the caller's contract). Randomness enters only through explicit uniform draws, so a fixed generator gives a fixed result.

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, V = 32000, K = 4 (and V = 128256, K = 5 for the published rows):

Over 600 rounds with a mismatched draft, the mean accepted length is 1.84 against the analytic sum_i acceptance^i of 1.86, where the per-token acceptance probability sum(min(p, q)) is 0.678.

Correctness

The accepted-length distribution and its mean track the analytic law of rejection sampling to the sampling floor at 600 rounds. Probabilities need not be normalized; masses are summed exactly. Greedy verification is deterministic and reproduces plain greedy target decoding exactly.

Requirements and limits

• target_probs [K+1, V] f32; for accept, draft_probs [K, V] warped identically to the target.
• x86-64 runs a scalar path; the O(V) reductions vectorize on aarch64.
• Distribution preservation depends on the caller warping both models' distributions identically.

References

Leviathan, Kalman, Matias, "Fast Inference from Transformers via Speculative Decoding" (2023); Chen et al., "Accelerating Large Language Model Decoding with Speculative Sampling" (2023).

License

Apache-2.0.