| --- |
| library_name: kernels |
| license: apache-2.0 |
| --- |
| |
| # spec-decode-ops |
|
|
| Fused sampling and speculative-decoding verification for autoregressive |
| decoding on NVIDIA Ampere and newer, loadable through `kernels`. The |
| reference baselines are the transformers processor chain (kept-set |
| equivalence exact on non-degenerate inputs) and the analytic laws of |
| rejection sampling, matched to a few parts in ten thousand. |
|
|
| Speculative decoding only pays off if verification is exact and cheap: a |
| draft model proposes K tokens, the target model scores them in one forward, |
| and the emitted stream must be distributed exactly as if the target had |
| decoded alone. This kernel implements the canonical rejection rule |
| (Leviathan et al. 2023; Chen et al. 2023) as one batched op, accepting each |
| draft token with probability `min(1, p/q)`, resampling the first rejection |
| from the renormalized residual `max(p - q, 0)`, and appending a bonus token |
| on full acceptance; the fused sampler covers the standard logits pipeline in |
| one kernel more. |
|
|
|  |
|
|
| *800 live verification rounds at K = 4 with a deliberately imperfect draft: |
| accepts green, first rejection red with its orange resample, bonus on full |
| acceptance. Mean accepted length 1.59 against the analytic law 1.55, with |
| the first-token distribution matching the target to the 800-round sampling |
| floor; 64 sequences verify at a 128,256-token vocabulary in 1.3 ms.* |
|
|
| ## Usage |
|
|
| ```python |
| import torch |
| from kernels import get_kernel |
| |
| sdo = get_kernel("phanerozoic/spec-decode-ops", version=1, trust_remote_code=True) |
| |
| logits = model(input_ids).logits[:, -1] # [B, V] |
| tok = sdo.sample(logits, temperature=0.8, top_k=50, top_p=0.95, |
| seed=1234, offset=step) # [B] int64 |
| |
| # draft model proposed draft_tokens [B, k]; both models were run over them |
| accept_len, out = sdo.verify(target_logits, # [B, k+1, V] |
| draft_logits, # [B, k, V] |
| draft_tokens, temperature=0.8, |
| seed=1234, offset=step) |
| ``` |
|
|
| `version` selects the release branch; `trust_remote_code` is required by |
| `kernels` for publishers without the trusted-publisher mark. Filtered |
| speculative decoding composes through `filter_logits`: apply the same filter |
| settings to both models' logits, then call `verify` with temperature 1. |
|
|
| ## API |
|
|
| | Symbol | Purpose | |
| |---|---| |
| | `sample(logits, temperature, top_k, top_p, min_p, repetition_penalty, prev_tokens, seed, offset)` | fused pipeline, one token per row | |
| | `filter_logits(...)` | the masked, temperature-scaled logits the eager processor chain would produce | |
| | `verify(target_logits, draft_logits, draft_tokens, temperature, seed, offset)` | rejection-sampling verification, exact target distribution | |
|
|
| ## Method |
|
|
| `sample` executes penalty, temperature, top-k, top-p, min-p, and the |
| categorical draw as one optional segmented sort plus one kernel; all three |
| filters are prefixes of the sorted order, so their cutoffs resolve in a |
| single walk, and the draw uses the Gumbel-argmax identity. Filter order and |
| boundary behavior match the transformers processors. Randomness is |
| counter-based Philox parameterized by `(seed, offset)`; results are bitwise |
| reproducible for fixed seed, offset, shape, and architecture. |
|
|
| ## Measured |
|
|
| On an L4, torch 2.12, `V = 50257` unless stated: |
|
|
| - Kept-set equivalence with the transformers chain is exact across |
| temperature/top-k/top-p/min-p/penalty combinations on non-degenerate |
| inputs; at a top-p boundary inside float32 cumsum precision the symmetric |
| difference carries < 1e-6 probability mass. |
| - Sampler distribution: TVD to the processor-chain analytic distribution |
| <= 1.9e-3 at 4.2e6 draws, zero mass on masked tokens. |
| - Verifier: emitted first-token distribution matches the target to TVD |
| 3.4e-3 at 2.1e6 trials with a mismatched draft; acceptance rate matches |
| `sum(min(p, q))` to 4e-4; the k=4 acceptance-length distribution matches |
| the analytic law to 3e-4. |
| - Determinism: identical `(seed, offset)` reproduce bitwise. |
| - Latency (`V = 128256`, median of 100): full filter pipeline 0.94 ms vs |
| 0.90/1.19 ms for the eager chain at M=1/M=8; the sort-free unfiltered |
| path runs 0.14 ms (6.4x/8.2x). |
|
|
| ## Requirements and limits |
|
|
| - NVIDIA GPU with compute capability 8.0+. |
| - logits in f32, bf16, or f16; `M * V < 2^31` per call. |
| - The filtered path is sort-dominated at M=1; a selection-based variant is |
| the planned successor. |
|
|
| ## References |
|
|
| Leviathan, Kalman, Matias, "Fast Inference from Transformers via Speculative |
| Decoding" (2023); Chen et al., "Accelerating Large Language Model Decoding |
| with Speculative Sampling" (2023); the transformers logits-processor chain. |
|
|
| ## License |
|
|
| Apache-2.0. |
|
|
