spec-decode-ops

Fused sampling and speculative-decoding verification for autoregressive decoding on NVIDIA Ampere and newer (sm_80 / sm_86 / sm_89 / sm_90).

sample executes the standard logits pipeline (repetition penalty -> temperature -> top-k -> top-p -> min-p -> 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, which requires neither normalization nor multinomial. The eager transformers chain performs the same pipeline as six or more full-vocabulary kernel passes. Greedy decoding and unfiltered sampling take sort-free single-kernel paths.

verify implements canonical speculative-decoding rejection sampling (Leviathan et al. 2023; Chen et al. 2023): draft token x_i is accepted with probability min(1, p_t(x_i)/p_d(x_i)); the first rejection is replaced by a sample from the renormalized residual max(p_t - p_d, 0); full acceptance appends a bonus token from the target's final position. The emitted stream is distributed exactly as the target model's distribution.

Randomness is counter-based Philox parameterized by (seed, offset); results are bitwise reproducible for fixed seed, offset, shape, and architecture.

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

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

Speculative verification:

# 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)
# out[b, :accept_len[b]+1] are the emitted tokens for sequence b

Filtered speculative decoding composes through filter_logits: apply the same filter settings to both models' logits, then call verify with temperature 1.

API

Function	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

Semantics

Filter order and boundary behavior match the transformers processors: penalty before temperature; top-p keeps the smallest sorted prefix whose cumulative probability reaches top_p, including the crossing token; min-p keeps p >= min_p * p_max; each filter keeps at least one token. top-p normalizes over the top-k prefix, matching warper application order.

Measured behavior

On an L4, torch 2.12, V = 50257 unless stated:

Kept-set equivalence with the transformers processor chain is exact across temperature/top-k/top-p/min-p/penalty combinations on non-degenerate inputs. At fully tied logits the kept sets are equal-sized subsets of the same exchangeable tie class (the reference's own subset is sort-order dependent); at a top-p boundary falling inside float32 cumsum precision the symmetric difference carries < 1e-6 probability mass.
Sampler distribution: total variation distance to the processor-chain analytic distribution <= 1.9e-3 at 4.2e6 draws across filter settings, with zero probability mass on masked tokens.
Verifier: emitted first-token distribution matches the target distribution to TVD 3.4e-3 at 2.1e6 trials with a deliberately mismatched draft; acceptance rate matches the analytic sum(min(p_t, p_d)) to 4e-4; the k=4 acceptance-length distribution matches the analytic law to 3e-4.
Determinism: identical (seed, offset) reproduce bitwise; changing offset changes draws.
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 (0.96x/1.27x); the sort-free unfiltered path runs 0.14 ms (6.4x/8.2x). The filtered path is sort-dominated at M=1; a selection-based variant is planned.