| --- |
| library_name: kernels |
| license: apache-2.0 |
| --- |
| |
| # 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](https://huggingface.co/kernels/phanerozoic/bitnet-cpu), |
| [quant-matmul](https://huggingface.co/kernels/phanerozoic/quant-matmul), |
| [gguf-gemv](https://huggingface.co/kernels/phanerozoic/gguf-gemv) |
| (linears), |
| [cpu-attn](https://huggingface.co/kernels/phanerozoic/cpu-attn), |
| [decode-ops](https://huggingface.co/kernels/phanerozoic/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 |
|
|
| ```python |
| 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 |
|
|
| | Symbol | Purpose | |
| |---|---| |
| | `verify_greedy(target_probs, draft_tokens)` | temperature-0 verify -> `(n_accepted, next_token)` | |
| | `accept(target_probs, draft_probs, draft_tokens, generator)` | speculative sampling -> `(n_accepted, next_token)` | |
| | `SpeculativeDecoder(draft_step, target_step, k, greedy, generator)` | draft/verify orchestration loop | |
|
|
| `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): |
|
|
| | op | Pi 5 | Pi 4 | |
| |---|---|---| |
| | `verify_greedy` | 33 us | 126 us | |
| | `accept` | 117 us | 657 us | |
| | `verify_greedy` (V=128256, K=5) | 59 us | 140 us | |
| | `accept` (V=128256, K=5) | 338 us | 754 us | |
|
|
| 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. |
|
|
