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

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.