I Built Speculative Decoding from a scratch— Here’s What Actually Broke

An honest account of building FlashSpec: an adaptive LLM inference engine with a Triton kernel and online bandit draft selection.

This is the story of what happened when I tried to build speculative decoding. This article is not a success story where everything worked perfectly. It’s a record of what I actually learned : including the mistakes, the silent bugs, and the parts that are still incomplete. I’m sharing it because I wish more people wrote about the messy middle of building real systems, not just the final polished results.

What FlashSpec Does

Speculative decoding is a technique that makes large language model inference faster without changing the final output distribution.

A small “draft” model quickly proposes several tokens at once. A larger “target” model then verifies all of them in a single forward pass. Tokens that pass the acceptance test are kept “for free.” Tokens that fail are replaced with the target model’s own output at that position.

The key property is that the final output distribution remains exactly the same as if you had sampled directly from the target model.

FlashSpec adds two main improvements:

1. GPU-Native Verification

Most implementations move the accept/reject decision to the CPU. This requires synchronizing data between the GPU and CPU at every decoding step, which adds overhead.

FlashSpec runs the entire verification step on the GPU using a custom Triton kernel. The kernel only reads two scalar values (log probabilities) per candidate token. This keeps memory usage constant, even if the vocabulary size is 32k or 128k tokens.

2. Online Bandit Draft Selection

Instead of choosing one draft model once and using it forever, FlashSpec treats draft model selection as a multi-armed bandit problem. It uses algorithms like UCB1 or Thompson sampling to adaptively pick which draft model to use during inference.

In theory, this allows the system to automatically find the best draft model for the current workload with bounded regret.

You can find the full code on GitHub, install it via PyPI, and explore a simple example notebook here. A preprint is also available on Zenodo.

Lesson 1: Writing the specification upfront forced me to answer questions I would normally postpone.

What should the exact mathematical formula be when a draft token is rejected? What numerical tolerance should the Triton kernel meet compared to a pure PyTorch implementation? At what sample size should the distribution-equivalence test run in CI?

These decisions are usually made reactively after bugs appear. Defining them in advance meant I could catch violations before those happen.

One clear example was the Kolmogorov-Smirnov (KS) test used to verify that FlashSpec’s output distribution matches the target model. The specification required running this test with 10,000 samples. In the initial code, it was only running with 1,000 samples.

The test passed at the lower sample size, so the problem stayed hidden. A KS test with only 1,000 samples has much less statistical power than one with 10,000 samples. Subtle distribution differences that should have been caught could have slipped through.

Fix: I increased the sample size to 10,000 and made the KS test a hard requirement in CI. If the test fails, the build fails.

Lesson: A specification that isn’t automatically checked will eventually be violated without anyone noticing.

Lesson 2: The Temperature Bug Was Invisible at Default Settings

This was the most educational bug in the project.

The mathematical rule is clear: when using temperature scaling, you must divide the raw logits before applying log_softmax. These two operations are not mathematically equivalent.

In the initial implementation, the rejection_sample() function accepted a temperature parameter, but it had no effect on the output. The log-probabilities it received were already computed inside another function (score_draft()), which applied log_softmax directly to the raw logits without any temperature scaling.

The parameter existed in the function signature and documentation. It was passed through the code. And it did nothing.

The bug was invisible in all tests that used the default value of temperature = 1.0 — which was every test. When temperature equals 1.0, dividing by 1.0 produces the same result as not dividing at all.

Fixing this required changing the architecture, not just editing one line:

# Before (temperature had no effect)
def score_draft(self, input_ids, draft_token_ids, gamma):
    logits = self._model(...).logits[..., -gamma:, :]
    return torch.log_softmax(logits.float(), dim=-1)

# After (temperature is applied before log_softmax)
def score_draft(self, input_ids, draft_token_ids, gamma, temperature=1.0):
    logits = self._model(...).logits[..., -gamma:, :]
    if temperature != 1.0:
        logits = logits / temperature          # ← Applied here
    return torch.log_softmax(logits.float(), dim=-1)

Three files needed to change. The temperature parameter was removed from rejection_sample() because that was the wrong place for it.

Lesson: In machine learning, wrong implementations often produce outputs that look correct: especially at default settings. Always write tests that check mathematical invariants at non-default values.

Lesson 3: Three Releases Shipped Before the Package Worked on Windows

After releasing version 0.1.0, testers immediately reported this error:

ERROR: Could not find a version that satisfies the requirement triton>=3.0.0

Triton only provides official wheels for Linux. There are no official Windows or macOS wheels. However, triton>=3.0.0 was listed as a required dependency in pyproject.toml. This made the package impossible to install on any non-Linux system from the first release.

Versions 0.1.0, 0.1.1, and 0.1.2 were all broken on Windows and macOS. All three versions are now yanked from PyPI.

The proper fix required changes in three places:

• Moving Triton to an optional gpu extra with a platform marker
• Adding graceful fallback code with a clear error message
• Pointing users to the pure-PyTorch reference implementation when Triton is unavailable

The package now installs cleanly on Windows, macOS, and Linux.

Lesson: Always test pip install your-package in a clean Windows environment before your first public release. It takes five minutes and catches an entire category of platform-specific problems.

Lesson 4: The Triton Kernel Was Slower Than PyTorch on the First Hardware I Tested

This result was disappointing, but important to document honestly.

On a Tesla T4 (Google Colab), the custom Triton verification kernel was significantly slower than the pure PyTorch reference at batch size 1: the most common case for single-user inference.

The reason is hardware-specific. The verification kernel is memory-bandwidth bound. The T4 has relatively low memory bandwidth compared to newer GPUs like the H100. On the T4, PyTorch’s highly optimized reference implementation was competitive. On higher-bandwidth hardware, the kernel’s much smaller memory footprint should provide a clearer advantage.

This result is documented in the README and JOSS paper with the hardware clearly stated. Performance claims for FlashSpec are conditional on higher-end GPUs. H100 benchmarking is still in progress.

Lesson: Writing a custom kernel does not automatically make something faster. Kernel performance depends heavily on the hardware. Always benchmark on your target hardware and clearly report which hardware was used.

Lesson 5: Property-Based Testing Found a Real Bug in Under a Minute

After adding Hypothesis (a property-based testing library) to the test suite, it quickly found a bug.

All my manually written tests used gamma=4 and batch_size=2,the shapes that felt natural while developing. Hypothesis generated gamma=1, batch_size=1 as one of its first test cases and triggered an out-of-bounds index error.

The shape was valid. I had simply never thought to test it.

I now run property-based tests on every CI build. These tests cover the full range of valid input values instead of just the few shapes I happened to write tests for.

Lesson: For any function where integer parameters affect shapes or indexing, add at least one property-based test. The effort is low, and the bugs it finds are often real.

Lesson 6: An Adaptive Algorithm Needs Real Variation to Be Useful

The theoretical results for the bandit looked good. Both UCB1 and Thompson sampling stayed well within their expected regret bounds in controlled experiments.

However, when I first tried the system with a real model (TinyLlama on a T4), there was only one draft model available. The bandit ran correctly but had nothing meaningful to adapt to.

The practical value of the bandit only appears when you have multiple draft models with genuinely different strengths — for example, a small fast drafter and a larger, more accurate one. Without real variation in the environment, the bandit becomes unnecessary overhead.

Lesson: An adaptive component is only useful if the environment contains meaningful variation worth adapting to. Validating an algorithm in isolation is necessary, but not sufficient.

Current State of the Project

What works today:

• The package installs cleanly on Windows, macOS, and Linux via pip install flashspec
• The optional gpu extra adds Triton kernels on Linux + CUDA
• The output distribution guarantee is enforced in CI (KS test at 10,000 samples)
• UCB1 and Thompson sampling satisfy their theoretical regret bounds
• First real measured performance: 44.2 tokens/second on TinyLlama-1.1B (4-bit) on a T4
• JOSS paper has been submitted

What is still in progress:

• Full H100 benchmarking with Llama-3–8B and Llama-3–70B (current headline numbers in the README are design targets)
• Minor code and lint fixes for readability

The core correctness, packaging, and distribution guarantee are solid. The main remaining work is validating performance on the target hardware.

One Thing I Would Do Differently

I would write the CI pipeline before writing the implementation, not alongside it.

The CI pipeline acts as the enforcement mechanism for the specification. When tests are added after the code already exists, there is a period where important invariants are not verified. Several issues in this project were only caught because a specification existed to compare against. They would have been caught much earlier if the CI gates had been running from the first commit.

The correct order is: Specification → CI → Implementation.

If you work on LLM inference systems and have thoughts on the kernel design or bandit formulation, I would genuinely value your feedback.

Try FlashSpec:

pip install flashspec            # Windows, macOS, Linux
pip install flashspec[gpu]       # Linux + CUDA (Triton kernels)

• GitHub: github.com/Mattral/FlashSpec
• PyPI: pypi.org/project/flashspec
• Example Notebook: Available in the repository under notebooks/