💻 computer science

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

FlashSpec is an open-source, adaptive speculative decoding engine that achieves exact target-model distribution preservation through a novel O(1) vocabulary-size Triton GPU kernel for on-device verification and an online bandit-based mechanism for dynamic draft-model selection, significantly accelerating LLM inference while eliminating CPU bottlenecks and manual tuning.

Original authors: Min Htet Myet

Published 2026-06-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Min Htet Myet

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a very long, complex story. You have a Master Author (the large AI model) who writes incredibly high-quality sentences but is very slow because they have to think through every single word carefully. You also have a Speedy Apprentice (the small draft model) who writes very fast but sometimes makes mistakes or uses the wrong tone.

Speculative Decoding is a technique where you let the Speedy Apprentice write a few words ahead of time, and then the Master Author quickly checks them. If the Master Author agrees, those words are accepted instantly. If not, the Master Author corrects them. This usually makes the whole process much faster.

However, the paper "FlashSpec" points out that existing systems have two major problems, like having a clumsy manager and a rigid hiring policy. Here is how FlashSpec fixes them:

1. The "CPU Bottleneck" Problem (The Clumsy Manager)

In current systems, every time the Master Author checks the Apprentice's words, the computer has to stop, send the data from the fast graphics card (GPU) to the slower main processor (CPU), wait for the CPU to do the math, and then send it back.

The Analogy: Imagine a race car driver (the GPU) who has to stop at every single mile marker to walk to a distant office (the CPU) to get a stamp of approval before continuing. This stops the car every time, killing the speed.
The FlashSpec Fix: The authors built a special, ultra-fast tool (a Triton kernel) that lets the Master Author check the Apprentice's words right there on the race car, without ever stopping to walk to the office. They only look at the two specific numbers needed for the check, ignoring the rest of the dictionary. This makes the check happen almost instantly, no matter how big the dictionary is.

2. The "Static Draft" Problem (The Rigid Hiring Policy)

Usually, you pick one Speedy Apprentice and stick with them for the entire job.

The Analogy: Imagine you hired an Apprentice who is great at writing poetry. But halfway through the job, you need to write a technical manual or a coding guide. Your poetry apprentice is now terrible at this, but you are forced to keep using them because you didn't plan to switch. You are wasting time.
The FlashSpec Fix: FlashSpec uses a "Smart Manager" based on Bandit Algorithms (a type of math used for making decisions under uncertainty).
- Instead of picking one apprentice forever, the system has a pool of different apprentices.
- At every step, the Smart Manager asks: "Who has been doing the best recently?"
- If the current apprentice starts making mistakes on a new topic, the manager instantly switches to a different apprentice who is better at that specific topic.
- It learns on the fly, without a human needing to tell it to switch.

The Results

The paper claims that by combining these two fixes:

Speed: The system runs much faster because it never stops to talk to the slow CPU.
Adaptability: It automatically switches strategies to stay fast, even when the topic changes from conversation to coding.
Accuracy: Despite all these shortcuts and switches, the final story is exactly the same as if the slow Master Author had written every word from scratch. The paper proves this mathematically and checks it with strict tests.

In short: FlashSpec is a new engine for AI that lets the AI "think ahead" without getting stuck in traffic, and it automatically hires the best "thought-leader" for whatever topic is being discussed at that exact moment. The authors have released this as open-source software so others can use it.

Technical Summary: FlashSpec

Problem Statement

Large Language Model (LLM) inference is fundamentally limited by memory bandwidth, as autoregressive generation requires loading all model parameters for every single token. Speculative decoding addresses this by using a smaller "draft" model to propose $\gamma$ candidate tokens, which a larger "target" model verifies in a single forward pass. While theoretically sound, existing open-source implementations suffer from two critical practical gaps:

CPU-Side Verification Bottleneck: The per-token accept/reject decision is typically executed in Python/PyTorch on the CPU. This necessitates a device-to-host synchronization at every decoding step, creating a pipeline stall proportional to the vocabulary size ( $V$ ), which ranges from 32k to 128k tokens in modern LLMs.
Static Draft Selection: Current systems select a draft model offline and keep it fixed throughout deployment. This ignores the non-stationary nature of inference workloads, where acceptance rates vary significantly across different prompt domains (e.g., conversational text vs. source code), leading to suboptimal throughput.

Methodology

FlashSpec is an open-source adaptive speculative-decoding engine designed to resolve both bottlenecks simultaneously while preserving the exact output distribution of the target model.

1. On-Device Verification Kernel

The core innovation is a Triton GPU kernel that performs the accept/reject test entirely on the device.

Mechanism: Instead of gathering full vocabulary logits, the kernel reads only two log-probability scalars per candidate token (one from the draft model, one from the target model).
Complexity: The kernel's SRAM footprint is $O(1)$ with respect to vocabulary size, independent of whether $V$ is 32k or 128k. Memory traffic per step is reduced from $O(B\gamma V)$ to $O(B\gamma)$ , where $B$ is the batch size.
Numerical Stability: The implementation uses log-space arithmetic to prevent underflow and correctly applies temperature scaling before the log-softmax operation, a common source of subtle bugs in prior implementations.

2. Online Bandit Draft Selection

FlashSpec frames the selection of the draft model as a $K$ -armed bandit problem.

Adaptation: At each decoding step, the system selects a draft model (an "arm") based on its recent performance. The reward is defined as the fraction of draft tokens accepted ( $\alpha$ ).
Algorithms: The system supports UCB1 (Upper Confidence Bound) and Thompson Sampling. These algorithms adapt to non-stationary acceptance rates without human retuning, minimizing cumulative regret over time.
Theoretical Guarantee: The bandit policy depends only on the history of past acceptances, not the current draft tokens or target distribution, ensuring the joint distribution of generated tokens remains identical to standard autoregressive sampling from the target model.

3. Correctness and Testing

The system includes a rigorous testing suite to ensure mathematical correctness:

Distribution Equivalence: A Kolmogorov–Smirnov (KS) test is enforced as a hard gate in the CI pipeline, comparing the speculative sampler's output against the target distribution over 10,000 samples at a significance level of $\alpha=0.01$ .
Property-Based Testing: The Hypothesis library is used to test edge cases (e.g., specific batch sizes and speculation lengths) that standard unit tests might miss.

Key Results

Initial measurements were conducted on a Tesla T4 GPU using TinyLlama-1.1B-Chat (4-bit NF4 quantization) with a speculation length of $\gamma=4$ .

Throughput: The system achieved 44.2 tokens/second with a mean acceptance rate ( $\bar{\alpha}$ ) of 0.75.
Latency: The p50 step latency was measured at 22.1 ms.
Kernel Efficiency: On the T4, the Triton kernel showed a 1.6 $\times$ speedup over a pure-PyTorch reference implementation at a batch size of 8, though it was slower at batch size 1 due to the T4's limited memory bandwidth (320 GB/s). The authors note that the kernel's benefits are expected to scale significantly on high-bandwidth hardware like the H100 (3.35 TB/s).
Bandit Performance: In a simulation with $T=10,000$ rounds and $K=3$ arms, the UCB1 selector achieved a cumulative regret of 100.2, well within the theoretical upper bound of $\approx 526$ . Thompson Sampling converged even faster with a regret of 18.9.

Significance and Claims

The paper positions FlashSpec as the first system to simultaneously address the CPU verification bottleneck and the static draft selection limitation in a provably correct, open-source framework.

Theoretical Contribution: The authors provide a formal proof (Theorem 1) that the output distribution of FlashSpec is identical to standard autoregressive sampling, regardless of the bandit's selection policy.
Practical Impact: By moving verification to the GPU and adapting draft selection online, FlashSpec aims to maximize throughput across diverse prompt distributions without manual intervention.
Engineering Rigor: The paper highlights specific engineering lessons learned, such as the necessity of applying temperature before log-softmax and the importance of large sample sizes (10,000) in distribution-equivalence testing to detect subtle shifts.

The authors note that full benchmarks on larger models (Llama-3-8B/70B) and comparisons with Medusa and EAGLE on H100 hardware are currently in progress and will be published in the associated repository. The software is released under the Apache 2.0 license.