FlashSpec: Adaptive Speculative Decoding with Online Bandit Draft Selection and Triton-Optimised Verification

Speculative decoding accelerates large language model (LLM) inference by having a small
draft model propose γ candidate tokens that a larger target model verifies in a single forward pass,
preserving the target model’s output distribution exactly. Existing systems have two practical
gaps: the per-token accept/reject decision is typically performed on the CPU—introducing a
pipeline stall proportional to vocabulary size—and the draft model is chosen offline, ignoring
shifts in prompt distribution across a deployment lifetime.

We present FlashSpec, an open-source adaptive speculative-decoding engine that addresses
both. Our first contribution is a Triton GPU kernel that performs the accept/reject test entirely
on-device, reading only two log-probability scalars per candidate token. The kernel’s SRAM
footprint is O(1) in vocabulary size, independent of whether the model has 32k or 128k tokens.
Our second contribution frames draft-model selection as a K-armed bandit problem and solves
it online with UCB1 [Auer et al., 2002] or Thompson sampling [Thompson, 1933, Chapelle and
Li, 2011], adapting to non-stationary acceptance rates without human retuning.
We prove that the output distribution of FlashSpec is identical to standard autoregressive
sampling from the target model (Theorem 1), and we enforce this guarantee continuously in CI
via a Kolmogorov–Smirnov test at significance α=0.01 over 10,000 samples.