FlashSpec: Adaptive LLM Inference with Speculative Decoding — Six Hard-Won Lessons

Speculative decoding accelerates LLM inference without changing the final output distribution. A smaller draft model proposes several tokens quickly; a larger target model verifies them in a single forward pass. Accepted tokens are kept; rejected tokens are replaced with the target model's output. The key constraint: the final output distribution must match direct sampling from the target model exactly.

FlashSpec's Two Core Innovations

1. GPU-Native Verification

Common implementations move accept/reject decisions to the CPU, requiring GPU-CPU synchronization each decoding step. FlashSpec keeps the entire verification step on the GPU using a custom Triton kernel. Each candidate token only reads two scalar log-probability values, so memory usage stays constant even with vocabularies of 32k or 128k tokens.

2. Online Bandit Draft Selection

Draft models are not fixed. FlashSpec models draft selection as a multi-armed bandit problem, using UCB1 or Thompson sampling to dynamically choose which draft model to use during inference. The theoretical goal: under bounded regret, let the system automatically find the best draft model for the current workload.

Lesson 1: Write Specifications Upfront

When a draft token is rejected, what is the exact mathematical formula? What numerical tolerance should the Triton kernel meet versus a pure PyTorch implementation? How many samples should the CI distribution-equivalence test use? These questions are easier to answer before bugs appear. Writing requirements into a spec exposes deviations earlier.

The Kolmogorov-Smirnov (KS) test verifies that FlashSpec's output distribution matches the target model. The spec required 10,000 samples, but the initial code only ran 1,000. Low sample count gave the test low statistical power, letting subtle distribution differences slip through. Fix: raise sample count to 10,000 and make the KS test a hard CI gate — build fails if the test fails. Lesson: a spec without automated checks will eventually be violated unnoticed.

Lesson 2: Temperature Bug Hidden by Default Settings

The most instructive bug in the project. With temperature scaling, the math is clear: raw logits must be divided by temperature before applying log_softmax; the two operations are not commutative. The original rejection_sample() accepted a temperature parameter, but the output was completely unaffected. The root cause was earlier: score_draft() computed log-probabilities by applying log_softmax directly to raw logits, with no temperature scaling. The parameter existed in the function signature and docs and propagated through the call chain, but was never used in the actual computation.

All tests used the default temperature = 1.0, so the bug remained invisible — dividing by 1.0 changes nothing. The fix required changes across three files: apply temperature to logits before log_softmax in score_draft(), and remove the parameter from rejection_sample() where it never belonged.

# 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 (apply temperature 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)

Lesson: incorrect ML implementations often produce "plausible" results; default parameters especially mask bugs. Mathematical invariant tests must cover non-default values.

Lesson 3: Three Releases Before Windows Install Worked

After the 0.1.0 release, testers immediately hit:

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

. Triton only provides official Linux wheels; no official Windows or macOS wheels exist. The pyproject.toml had declared triton>=3.0.0 as a required dependency, making the package uninstallable on non-Linux systems from day one. Versions 0.1.0, 0.1.1, and 0.1.2 were all unusable on Windows and macOS and have been yanked from PyPI.

The full fix involved three changes:

Move Triton to an optional gpu extra with platform markers.

Add graceful fallback code and clear error messages.

When Triton is unavailable, guide users to the pure PyTorch reference implementation.

Now Windows, macOS, and Linux can all install the package. Lesson: before the first public release, run pip install your-package in a clean Windows environment — five minutes catches a whole class of platform issues.

Lesson 4: Triton Kernel Slower Than PyTorch on T4

A disappointing but honest result. On a Tesla T4 (Google Colab) with batch size 1 — the most common single-user inference scenario — the custom Triton verification kernel was noticeably slower than the pure PyTorch reference implementation. The verification kernel is memory-bandwidth bound; the T4's bandwidth is lower than newer GPUs like the H100, and PyTorch's highly optimized reference implementation remains competitive on T4. On higher-bandwidth hardware, the Triton kernel's smaller memory footprint should show its advantage. Both the README and the JOSS paper record these results and explicitly state the test hardware. FlashSpec's performance claims assume higher-end GPUs; H100 benchmarks are ongoing. Lesson: custom kernels do not guarantee speedups. Performance is highly hardware-dependent; benchmarks must run on target hardware and the test platform must be documented.

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

Hypothesis, a property-based testing library, was added to the test suite and quickly uncovered a previously missed issue. Hand-written tests always used gamma=4 and batch_size=2 because those shapes were convenient during development. Hypothesis early generated gamma=1, batch_size=1, which immediately triggered an out-of-bounds array error. This shape is perfectly valid but had never been considered. Now every CI run executes property-based tests covering the full valid input range, not just a few hand-picked shapes. Lesson: any integer parameter that affects shape or indexing deserves at least one property-based test. Low investment, often catches real bugs.

Lesson 6: Adaptive Algorithms Need Real Differences to Be Valuable

The bandit theory checks out: in controlled experiments, UCB1 and Thompson sampling both stay within their expected regret bounds. But when the system was first connected to real models, the environment gave the bandit no meaningful choice space: running TinyLlama on a T4 with only one draft model available. The algorithm executed correctly, but there was nothing to adapt to. Bandit value depends on having multiple draft models with different trade-offs — e.g., a small fast drafter and a larger more accurate one. Without real differences in the environment, adaptive selection only adds overhead. Lesson: an adaptive component's precondition is that the environment actually contains variations worth adapting to. Verifying the algorithm in isolation is necessary but not sufficient.

Current Project Status

Working now: pip install flashspec installs on Windows, macOS, and Linux.

Optional gpu extra adds Triton kernels on Linux + CUDA.

CI enforces output distribution guarantee (10,000-sample KS test).

UCB1 and Thompson sampling satisfy their theoretical regret bounds.

First measured performance: 44.2 tokens/second running TinyLlama-1.1B (4-bit) on T4.

JOSS paper submitted.

Ongoing work:

Full H100 benchmarks with Llama-3-8B and Llama-3-70B (current README numbers are design targets).

Minor code and lint fixes for readability.

Core correctness, packaging, and distribution guarantees are stable. Main remaining work is performance validation on target hardware.

If Starting Over

The CI pipeline should be written before implementation, not in parallel. CI is the enforcement mechanism that makes a spec real. Writing code first, adding tests later, leaves a window where critical invariants go unverified. Several issues in this project were caught only because a spec existed to check against; with a CI gate from the first commit, they would have surfaced earlier. Correct order: Specification → CI → Implementation.

Try FlashSpec:

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

GitHub: github.com/Mattral/FlashSpec by Min Htet Myet