⚡ FlashSpec

Lossless speculative decoding for HuggingFace causal language models
Algorithm 1 rejection sampling and exact greedy · static KV cache + CUDA graphs on CUDA · experimental bandit and kernels

Project status: alpha, version 0.2.0.

What was measured, and only this: one model pair (Qwen/Qwen2.5-1.5B-Instruct verified by Qwen/Qwen2.5-0.5B-Instruct), one Tesla T4, fp16, batch 1, 10 MT-Bench and 10 HumanEval prompts, run from main at git 8d036d7 with the default FlashSpecConfig(compile="auto"), which turned the compiled path (static KV cache and CUDA graphs) on. On that setup FlashSpec was faster than HuggingFace generate compiled the same way in every cell of the table below: by at least 1.3x on both HumanEval cells, and by less on MT-Bench. With the eager path (compile=False) it was slower than plain HuggingFace generate in 3 of the 4 cells of the same run (generated table further down). No other GPU, model pair or batch size has been measured. Every H100 number in the docs is a design target.

• The first call compiles. On CUDA with Triton, the default compile="auto" compiles the drafter and target forwards and captures CUDA graphs on the first generate, which can take minutes. Later calls reuse them. The tables above and below exclude that first call (warmup prompts absorb it). To skip compilation, pass FlashSpecConfig(compile=False) (the eager path, slower here) or --compile off to the server. The engine logs one line saying whether the compiled path is on and why.

• Correctness is checked in CI on CPU: a chi-squared test that emitted tokens follow the target when draft and target differ, and a token-for-token comparison of greedy output with generate(do_sample=False). On the GPU in fp16, kernels for different shapes can round differently, so each greedy GPU run reports how many prompts matched (last columns of the table).

• The bandit is experimental. Every measured run used one drafter, so the selector has never been measured choosing between drafters.

• The Triton and assembly kernels are experimental and not on the generation critical path; see Experimental kernels.

• sampling.variant="typical" is lossy; "rejection" and "greedy" are the lossless variants.

This is a single-device library plus a small HTTP server, not a production serving stack.

Measured at git 8d036d7: Qwen/Qwen2.5-1.5B-Instruct verified by Qwen/Qwen2.5-0.5B-Instruct, Tesla T4, fp16, batch 1, gamma 4, 10 prompts per set (chat template), 128 new tokens. One model pair, one GPU. Speedups are medians of per-prompt ratios within the same run. The headline baseline is HuggingFace generate with a static cache compiled the same way (CUDA graphs); eager HuggingFace and assisted generation are context. Cells are the stored JSON values.

Greedy: compiled FlashSpec produced the same tokens as hf_greedy_compiled, token for token (per-prompt output hashes are in the JSON), on 10/10 humaneval prompts and 10/10 mt_bench prompts. Cold start: not recorded in this bundle. The first compiled call compiles the model and captures CUDA graphs, which can take minutes; the timed numbers above exclude it (warmup prompts absorb it).

What is FlashSpec?

Speculative decoding lets a small "draft" model propose γ tokens that a larger "target" model checks in one forward pass. FlashSpec's batch-1 step is:

1. the drafter proposes γ tokens (γ incremental forwards on its KV cache);
2. the target runs one forward over the pending token(s) plus the γ drafts;
3. the accept/reject test runs on the device in PyTorch: Leviathan et al. Algorithm 1 rejection sampling at temperature > 0, or the longest prefix that matches the target's argmax at temperature 0 (greedy);
4. one device-to-host read per step returns the accept count and the next token, and both caches are cropped by integer bookkeeping.

On CUDA with Triton available, FlashSpecConfig.compile runs that loop on static KV caches and replays the fixed-shape drafter decode and target verify calls as CUDA graphs (torch.compile(mode="reduce-overhead")).

Two parts are experimental. The UCB1/Thompson selector can choose between registered drafters, but it has only been measured with one. The Triton verification kernels and the optional assembly residual are not called on the generation path (see Experimental kernels).

On finite probabilities the output distribution follows Leviathan et al., 2023, Theorem 1. The CI check is a chi-squared test on discrete token ids. tests/integration/test_output_distribution.py uses fixed pairs with draft distribution q different from target distribution p, at gamma 1 and gamma 4. It compares the first emitted token to p, and, when gamma is 4, the second emitted token as well. N is 30,000 and the vocabulary is 6, so each expected count stays above 5. The critical value is the chi-square 1% point for 5 degrees of freedom, 15.086. The same check runs on the PyTorch residual and on a kernel that implements the same formula. An engine test uses a toy target whose next token depends on the previous token and compares the first two generated tokens with sampling from that target. tests/integration/test_e2e_sampling.py is the gamma 1 case where the draft equals the target, so every draft token is accepted. It does not exercise the residual. A failing chi-squared test fails the suite.

Quickstart

git clone https://github.com/Mattral/FlashSpec && cd FlashSpec

# Linux + CUDA (accelerated Triton kernels):
pip install -e ".[dev,gpu]"

# Windows / macOS (pure-PyTorch reference kernels, numerically identical):
pip install -e ".[dev]"

# Optional CPU residual kernel (Linux x86-64, flashspec-asm-kernel>=0.1.2;
# other platforms install nothing):
pip install -e ".[asm]"

# Optional single-device HTTP server:
pip install -e ".[serve]"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from flashspec import FlashSpecConfig, BanditConfig, SamplingConfig, TRITON_AVAILABLE
from flashspec.bandit import UCB1Selector
from flashspec.engine import SpeculativeEngine, TargetModel, load_drafters

print(f"Triton kernels available: {TRITON_AVAILABLE}")

# Public checkpoints. Both report vocab_size 32000.
# Llama-3 reports vocab_size 128256. Do not pair it with this drafter.
drafter_id = "Felladrin/Llama-160M-Chat-v1"
target_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

config = FlashSpecConfig(
    device="cpu",
    dtype="float32",
    drafter_name=drafter_id,
    target_name=target_id,
    bandit=BanditConfig(n_arms=1, strategy="ucb1"),
    sampling=SamplingConfig(gamma=4, temperature=1.0),
    max_new_tokens=16,
)
# TargetModel(hub_id, device) passes device_map and needs accelerate,
# which this package does not depend on. Load the weights here and inject
# them, which is what benchmarks/harness.py does.
drafter_module = AutoModelForCausalLM.from_pretrained(
    drafter_id, torch_dtype=torch.float32
).to(config.device)
target_module = AutoModelForCausalLM.from_pretrained(
    target_id, torch_dtype=torch.float32
).to(config.device)
drafter_module.eval()
target_module.eval()
drafters = load_drafters(
    [drafter_id],
    device=config.device,
    dtype=torch.float32,
    temperature=config.sampling.temperature,
    models={drafter_id: drafter_module},
)
target = TargetModel(
    target_id,
    torch.device(config.device),
    torch.float32,
    model=target_module,
)
engine = SpeculativeEngine(
    config,
    drafters,
    target,
    UCB1Selector(n_arms=1, window_size=config.bandit.window_size),
)
tokenizer = AutoTokenizer.from_pretrained(target_id)
prompt = tokenizer("The capital of France is", return_tensors="pt").input_ids
result = engine.generate(prompt, max_new_tokens=16)
new_ids = result.output_ids[0]
new_ids = new_ids[new_ids != -1]
print(tokenizer.decode(new_ids, skip_special_tokens=True))

Both checkpoints are public. benchmarks/harness.py refuses a drafter and a target whose vocabulary sizes differ and writes no result file. SpeculativeEngine itself does not apply that check. notebooks/01_quickstart.py is a separate CPU toy: the names are "toy", the log-probs are random, and it does not download these checkpoints. With one drafter the bandit cannot adapt, and the loader says so.

Results

The generated GPU blocks show the stored JSON values. The hand-written tables below them use three significant figures; the JSON and CSV files keep full precision. Each figure is classified in benchmarks/results/PROVENANCE.md. Commands and file names are in benchmarks/results/README.md.

Measured — GPU, real prompts (MT-Bench and HumanEval, chat template)

This block is generated by scripts/ingest_gpu_bundle.py --readme-commit 8d036d7 from the result JSONs of the 0.2.0 bundle only (benchmarks/results/gpu/8d036d7/, run from main). Do not edit it by hand. The earlier feature-branch bundles (36134c0, the Phase 4 eager run that also covered TinyLlama, and 52993e4, the first compiled run) stay under benchmarks/results/gpu/ as history and are not shown here.

Qwen/Qwen2.5-1.5B-Instruct verified, Qwen/Qwen2.5-0.5B-Instruct drafting

GPU Tesla T4 · torch 2.11.0+cu130 · gamma 4 · 20 result files · git 8d036d7. Prompts: chat_template; n per row: 10. Cells are the stored JSON values, not rounded.

• forward_cost (humaneval): c eager 0.8840264407247382, c compiled 0.47651099333876806
• forward_cost (humaneval): c eager 0.8922046962713643, c compiled 0.4770852877691722
• forward_cost (mt_bench): c eager 0.8978235042101232, c compiled 0.47562151975963995
• forward_cost (mt_bench): c eager 0.9040499147358744, c compiled 0.47418886402988414

Speedups are medians of per-prompt ratios against the method in the same run. NaN/inf counts come from one untimed probe generation per method on the first prompt.

Measured, 0.1.5 engine — TinyLlama-1.1B-Chat float32 · CPU · γ=4 · git f59b5b4

One HumanEval prompt, 16 new tokens, one drafter. Speedup is below 1.

Measured, 0.1.5 engine — TinyLlama-1.1B-Chat float16 · Tesla T4 · γ=4 · git cb02ca1

Eight synthetic prompts, one drafter. Speedup is below 1.

Microbenchmarks and historical

These rows are a random network, a sampling kernel, or the historical NF4 prompt. They sit below the paired end-to-end runs. The kernel and bandit measurements are under Experimental components.

CPU toy and residual — git da8182b

Random in-process network, and a PyTorch residual microbench. The toy speedup is below 1. The bandit row is a synthetic sampler, not that network.

An older notebook cell assigns 44.2 tok/s, acceptance 0.75, p50 22.1 ms; no timing log exists, so it is not used as a result.

Gamma sweep — measured (toy logprobs, synthetic)

Acceptance rate and throughput across speculation lengths γ ∈ {1, 2, 4, 8, 16}. Data in benchmarks/results/gamma_sweep.csv.

Throughput here is the sampling kernel on random log-probs. The CSV column measurement_kind is sampling_kernel_only_random_logprobs. It is not end-to-end inference and it is not an H100 result.

No H100 comparison exists, and Medusa and EAGLE are not implemented here, so no comparison with them exists.

Experimental components

Bandit (experimental, single-arm measured only)

Every language-model run in this README used one drafter, so the selector has not been measured choosing between drafters. The regret numbers below come from a synthetic K=3 simulation.

Bandit regret — measured, synthetic (T=10,000, K=3 arms)

In this simulation UCB1 and Thompson sampling stay inside the O(√(KT log T)) bound over 10,000 rounds. UCB1 cumulative regret at T=10,000: 100.2 (theory bound: ~526). Thompson cumulative regret at T=10,000: 18.9.

The figure is paper/figures/bandit_regret.jpg. The saved output is notebooks/02_bandit_analysis.ipynb cell 3 (0-based index, execution count 4): UCB1 100.2, Thompson 18.9, Oracle 0.0. This is a synthetic K=3 arm experiment, not a throughput gain on a language model.

Experimental kernels (not on the critical path)

The engine does not call these kernels. A speculative step's time is the drafter and target forwards (milliseconds); the accept/reject test and the residual are microseconds and run in PyTorch on the device.

• Triton verify_tokens / gather_accepted (flashspec.kernels). verify_tokens tiles over (batch × γ) positions; each thread reads two scalars (log p, log q at the draft token index) and writes one bool and one int32, so SRAM use does not grow with the vocabulary. It matches the pure-PyTorch reference numerically. There is no GPU CI runner: notebooks/colab_triton_tests.py runs the @pytest.mark.gpu parity tests on Colab and writes colab_triton_tests.log.
• Assembly residual (flashspec-asm-kernel, the [asm] extra, Linux x86-64). A residual evaluation is microseconds against millisecond model forwards, so it does not change end-to-end throughput.

Triton verification kernel vs the PyTorch reference — measured (Tesla T4)

Saved output of notebooks/03_kernel_profiling.ipynb cell 3 (0-based index, execution count 4), printed there as a Tesla T4 run. At batch 1 the Triton kernel is slower than the PyTorch reference.

The Triton kernel is slower than the pure-PyTorch reference at batch 1 on this T4 measurement. A faster H100 result has not been measured. Do not treat the H100 roofline in docs/kernels.md as a profile.

Assembly residual ([asm] extra) — Linux, git 4e48c55, then git 1b0cda3

Milliseconds, backend asm. Compare backends inside one run. The two runs do not repeat the same ordering.

How it works

Batch-1 fast path (variant "rejection" or "greedy"):

sequenceDiagram
    participant E as SpeculativeEngine
    participant B as Bandit (experimental)
    participant D as Drafter
    participant T as Target

    loop every decoding step
        E->>B: select() → arm index
        E->>D: draft γ tokens (γ forwards on its KV cache)
        E->>T: verify pending + γ tokens [one forward]
        E->>E: rejection_step or greedy_step on the device (PyTorch)
        E->>E: one host read: n_accepted, next token
        E->>D: commit: crop the cache, set the pending token
        E->>T: commit: crop the cache, set the pending token
        E->>B: update(arm, emitted tokens / step seconds)
    end

Batching. A batch of more than one row uses the generic loop. Rows that accept different numbers of tokens finish the step at different context lengths. Each loop forwards only the largest group of rows that still share a length. Other lengths wait for a later loop. The KV cache is rolled only when that group has a single row. A larger group full-forwards. The returned tensor fills a shorter row with -1 so the tensor is rectangular. That fill is not an attention mask. The engine does not pad unequal prompts into one masked batch.

Correctness. The accept/reject criterion is u_i < min(1, p(xᵢ)/q(xᵢ)) computed in log-space to avoid underflow. Rejected positions sample a residual token from max(0, p−q) / ‖max(0, p−q)‖₁ (Algorithm 1, Leviathan et al. 2023). Together these preserve the target distribution exactly.

Compiled path. With FlashSpecConfig.compile on a CUDA device, the drafter and target run on preallocated static KV caches; a crop moves the write position instead of slicing tensors, and the fixed-shape decode and verify calls replay as CUDA graphs. The prompt prefill stays eager.

Bandit (experimental). UCB1 score for arm k at round t: μ̂ₖ + c · √(2 log t / nₖ). Unpulled arms are selected first. The reward is emitted tokens per step second. In a synthetic K=3 simulation UCB1 regret at T=10,000 is 100.2 (theory bound ≈ 526) and Thompson's is 18.9. No language-model run has used more than one drafter.

See docs/architecture.md for the full component diagram and proof sketch.

Installation

# PyPI — cross-platform (Windows/macOS/Linux):
pip install flashspec

# PyPI — adds Triton (Linux only): the experimental Triton kernels.
# Linux CUDA builds of PyTorch already ship Triton, which torch.compile
# needs for the compiled path.
pip install flashspec[gpu]

# Optional AVX2 residual kernel. The wheel is a Linux x86-64 ELF library.
# On Windows and macOS the extra installs nothing and the PyTorch residual stays selected.
pip install flashspec[asm]

# Optional single-device HTTP API (FastAPI). One GPU or CPU, no streaming.
pip install flashspec[serve]

# From source:
git clone https://github.com/Mattral/FlashSpec && cd FlashSpec
pip install -e ".[dev,gpu]"   # omit ",gpu" on Windows/macOS

# Docker (Linux + CUDA). No image is published; build it locally:
docker build -f deploy/Dockerfile -t flashspec .
docker run --gpus all flashspec make test

Windows / macOS: pip install flashspec works out of the box. Triton has no official PyPI wheels outside Linux, so the compiled path and the experimental Triton kernels are Linux + CUDA only. The engine's accept/reject test is PyTorch and runs everywhere PyTorch does. The kernels' pure-PyTorch reference is flashspec.kernels._reference. Check availability at runtime: from flashspec import TRITON_AVAILABLE.

Requirements

The asm extra installs flashspec-asm-kernel >= 0.1.2 and < 2 on Linux x86-64. Release 0.1.2 is a manylinux wheel that contains libflashspec_asm.so. Release 0.1.0 was an sdist with no shared library. An import error, or an OSError from abi_version(), keeps the pure-PyTorch residual. A residual evaluation is microseconds, and a model forward is milliseconds, so this kernel does not change end-to-end throughput. benchmarks/results/residual_linux_gha.json and benchmarks/results/residual_linux_gha_1b0cda3.json are source builds from the flashspec-asm repo, not a timing of the 0.1.2 wheel. The PyPI project flashspec at version 0.1.2 was yanked by the maintainer. Do not install flashspec==0.1.2. That yank does not apply to flashspec-asm-kernel.

Testing

make test           # unit + integration on CPU (no GPU required, ~2 min)
make test-chaos     # adversarial bandit tests
make test-gpu       # pytest -m gpu; needs CUDA. This repository has no GPU runner.
make bench-quick    # sampling-kernel smoke test, no model weights, writes no result file
make bench-cpu      # residual microbench, synthetic bandit demo, tiny_cpu.yaml
make bench          # runs every config in benchmarks/configs; skips any checkpoint this machine cannot load

The distribution check in the suite is a chi-squared test (tests/integration/test_output_distribution.py, N = 30,000, vocabulary 6, critical value 15.086 at 5 degrees of freedom and α = 0.01). It is not a Kolmogorov–Smirnov test. The CPU test job runs it. Pytest fails under 83% line coverage, the measured unit-and-integration coverage rounded down. The CI lint job also runs mypy --strict flashspec.

Repository layout

flashspec/          Installable package (engine, kernels, bandit, sampling, metrics)
benchmarks/         Runners, configs, and results/ for all paper numbers
tests/              Unit · integration · chaos suites (83% coverage floor)
notebooks/          Quickstart · bandit regret analysis · kernel profiling
docs/               MkDocs site (architecture, kernels, bandit, benchmarks)
paper/              Draft Zenodo preprint (article class) and bibliography
paper/joss/         JOSS-format draft, not under review
paper/figures/      Vector and raster figures (bandit_regret.jpg generated)
scripts/            CLI tools (export draft, download models, profile kernels)
deploy/             Dockerfile · docker-compose.yml · k8s manifests

Links

Citation

@misc{myet2026flashspec,
  title = {{FlashSpec}: Adaptive Speculative Decoding with Online Bandit
             Draft Selection and {Triton}-Optimised Verification},
  author = {Myet, Min Htet},
  year = {2026},
  month = jun,
  howpublished = {Zenodo preprint},
  doi = {10.5281/zenodo.20766119},
  url = {https://doi.org/10.5281/zenodo.20766119},
}

Preprint DOI: 10.5281/zenodo.20766119 (June 2026). That record is the citable archive. paper/flashspec.tex is a later draft (October 2026, software 0.2.0) and has not been deposited.

Machine-readable metadata in CITATION.cff.

License

Apache 2.0 — see LICENSE.