💬 NLP

FlexDraft: Flexible Speculative Decoding via Attention Tuning and Bonus-Guided Calibration

FlexDraft is a lossless speculative decoding framework that accelerates LLM inference across varying batch sizes by combining attention tuning for high-quality block diffusion drafting, bonus-guided calibration to resolve draft verification mismatches, and a dynamic switching mechanism that optimizes between parallel and sequential execution modes.

Original authors: Yaojie Zhang, Jianuo Huang, Junlong Ke, Yuhang Han, Yongji Long, Tianchen Zhao, Biqing Qi, Linfeng Zhang

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yaojie Zhang, Jianuo Huang, Junlong Ke, Yuhang Han, Yongji Long, Tianchen Zhao, Biqing Qi, Linfeng Zhang

Original paper licensed under CC BY 4.0 (http://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 long story, but you have a very strict editor (the Target Model) who is incredibly smart but moves very slowly because they have to check every single word before writing the next one. This makes writing the story take forever.

To speed things up, you hire a fast, energetic assistant (the Drafter) to guess the next few words for you. The assistant writes a whole paragraph ahead of time, and then the slow editor quickly checks if those guesses were right. If they are right, you get to skip the editor's slow thinking time for those words. This is called Speculative Decoding.

However, the old way of doing this had two big problems:

The Waiting Game: The assistant would write a guess, then wait for the editor to finish checking it before writing the next guess. They were taking turns, which wasted time.
The "What If" Confusion: In newer, faster methods, the assistant tries to write multiple different future paragraphs at once, hoping to cover all possibilities. But the assistant doesn't know which paragraph the editor will actually approve. This leads to the assistant wasting energy writing paragraphs that get thrown away, and the editor having to check too many options, which slows everything down again when you have many stories to write at once (large batch sizes).

Enter FlexDraft: The Smart, Flexible Assistant

The paper introduces FlexDraft, a new system that fixes these problems so you can write faster without losing quality. Here is how it works, using simple analogies:

1. The "Specialized Glasses" (Attention Tuning)

Usually, to make the assistant better at guessing, you have to retrain their entire brain, which is expensive and risky (it might make them forget how to speak properly).

FlexDraft is different. It puts a pair of "specialized glasses" only on the assistant's eyes for the very last few steps of thinking.

How it works: The assistant keeps all their original knowledge (the frozen brain) but learns to use these new glasses to look at "blank spaces" (mask tokens) and fill them in all at once.
The Benefit: The assistant becomes a fast, parallel guesser without needing a full brain transplant. It stays true to the editor's style, so the guesses are high-quality and safe.

2. The "Magic Note" (Bonus-Guided Calibration)

Here is the tricky part: When the editor checks the assistant's guesses, sometimes the editor says, "No, that word is wrong, but here is the correct word to start with instead." This correct word is called the Bonus Token.

In the old "fast" methods, the assistant had to guess the next paragraph before seeing this Magic Note. So, the assistant might guess a story about a pirate, while the editor wanted a story about a chef. The assistant's guess was wrong because they didn't have the note.

FlexDraft adds a Magic Note system:

How it works: Once the editor reveals the Bonus Token, FlexDraft uses a tiny, fast calculator to instantly adjust the assistant's previous guesses to match that new direction.
The Benefit: Even if the assistant guessed blindly at first, the system quickly "steers" the guess to align with the editor's correction. This means fewer guesses get rejected.

3. The "Traffic Cop" (Flex Decoding)

The biggest problem with the old "fast" methods was that they tried to write every possible future story at once. If you only have one story to write (small batch), this is fine. But if you have 16 stories to write at once (large batch), the assistant gets overwhelmed trying to write 16 different futures for every single story, and the editor gets bogged down checking them all.

FlexDraft acts like a smart Traffic Cop:

Small Traffic (Small Batch): When there are few stories, the cop says, "Go! Write all possible futures at the same time!" This overlaps the writing and checking to save time.
Heavy Traffic (Large Batch): When there are many stories, the cop says, "Stop! Don't write all the possibilities. Just write the one most likely to be approved."
The Benefit: It switches strategies automatically. It avoids the "wasted energy" of writing too many options when the system is busy, preventing the speed from crashing.

The Result

By combining these three tricks, FlexDraft allows the slow editor to process text much faster.

No Quality Loss: The final story is exactly the same as if the slow editor had written it word-by-word.
Huge Speedup: On tests with a popular AI model (Qwen3-8B), FlexDraft made the system 4.59 times faster than the standard slow method.
Scalability: It works well whether you are writing one story or managing a huge crowd of stories, unlike previous methods that got slow and clunky when the crowd got big.

In short, FlexDraft is a flexible, lossless way to let a fast assistant help a slow editor, ensuring they work together smoothly without wasting time or making mistakes.

Technical Summary: FlexDraft

Problem Statement

Speculative decoding accelerates Large Language Model (LLM) inference by using a lightweight drafter to propose multiple candidate tokens for parallel verification by a target model. While effective, conventional sequential speculative decoding suffers from mutual waiting between drafting and verification phases, as well as high memory access overhead due to repeated state exchanges.

Recent parallel speculative decoding methods attempt to overlap drafting and verification within a single target forward pass to mitigate mutual waiting. However, this paradigm introduces two fundamental limitations that degrade performance, particularly at large batch sizes:

Bonus Token Uncertainty: In parallel decoding, draft tokens are generated before the verification outcome (specifically the "bonus token" sampled at the first rejection point) is known. This leads to a mismatch between the draft distribution and the verification constraints, reducing acceptance rates.
Acceptance Length Uncertainty: Since the number of accepted tokens is unknown during drafting, parallel methods must prepare candidate branches for all possible accepted prefixes. This forces the target model to process a rapidly growing number of "mask tokens" (future drafts) alongside verification tokens, causing computational overhead to scale quadratically ( $O(N^2)$ ) with the draft block length. This redundancy causes throughput gains to collapse at larger batch sizes.

Methodology: FlexDraft

FlexDraft is a lossless speculative decoding framework designed to adapt flexibly to varying batch sizes while preserving the target model's distribution. It addresses the aforementioned limitations through three core components:

1. Attention Tuning (Attn Tuning)

To enable block diffusion drafting (generating multiple tokens in parallel) without full model retraining or quality degradation:

Architecture: The framework keeps the target model's autoregressive path (FFN layers and original attention projectors) frozen.
Trainable Parameters: Only the attention projectors ( $Q, K, V, O$ ) for the final few layers are tuned specifically for mask tokens.
Mechanism: Mask tokens are appended to the clean prefix. In deep layers, these mask tokens attend bidirectionally within their block and to the clean prefix, while remaining isolated from other blocks. This allows the drafter to inherit the target model's pretrained knowledge via frozen FFNs while learning to predict multiple future tokens in parallel using a minimal number of additional parameters (approx. 6% of the target model size).

2. Bonus-guided Calibration

To resolve the mismatch caused by bonus token uncertainty:

Mechanism: A lightweight Multi-Layer Perceptron (MLP) is applied after the Language Modeling (LM) head.
Input: The MLP takes the concatenation of the resolved bonus token embedding and the hidden state of the current mask position.
Output: It produces an additive calibration bias to the draft logits. This steers the draft distribution to be consistent with the bonus token determined by the verifier, effectively aligning the draft with the verification constraints even though the draft was generated without knowledge of the bonus token.

3. Flex Decoding

To handle the trade-off between mutual waiting (small batches) and redundant computation (large batches):

Dynamic Switching: The system dynamically switches between Parallel Draft & Verify (for small batch sizes) and Sequential Draft then Verify (for large batch sizes) based on a batch size threshold.
Selective Verification (Small Batches): In parallel mode, the system estimates the probability of accepted lengths using cumulative draft confidence. It prunes low-confidence candidate branches, verifying only a subset of high-confidence prefixes to reduce target forward overhead without changing the draft block length.
Decoupled Execution (Large Batches): In sequential mode, once the accepted length is resolved, the drafter generates only the single branch conditioned on the verified prefix. This reuses the KV cache from the target verification step, eliminating the redundant branches that cause $O(N^2)$ overhead.

Key Contributions

Lossless Block Diffusion Drafting: A method to enable parallel token generation by tuning only final-layer attention projectors, preserving the target distribution and avoiding costly continual pretraining.
Bonus-guided Calibration: A mechanism to mitigate draft-verification mismatch by calibrating draft logits based on the resolved bonus token, significantly improving alignment.
Batch-Adaptive Execution: A unified framework that switches between parallel and sequential strategies and employs confidence-based pruning to prevent speedup collapse at large batch sizes.

Experimental Results

Experiments were conducted on Qwen3 models (1.7B, 4B, and 8B) across benchmarks including GSM8K, MATH, HumanEval, MBPP, and MT-Bench.

Performance: FlexDraft achieves an average 4.59× speedup over autoregressive decoding on Qwen3-8B without quality loss.
Comparison: It outperforms strong baselines including EAGLE-3, DFlash, BiTA, and Apple MTP. Specifically, on Qwen3-8B, FlexDraft achieves a 5.88× speedup on GSM8K and 5.79× on MATH, surpassing DFlash (4.55× and 5.50× respectively).
Batch Size Robustness: Unlike existing parallel methods that suffer from speedup collapse at large batch sizes due to redundant computation, FlexDraft maintains superior speedup across batch sizes (1 to 16) by switching to sequential execution and pruning low-confidence branches when necessary.
Ablation Studies:
- Attn Tuning alone provides significant gains by reusing target knowledge.
- Bonus-guided Calibration further improves speedup by aligning drafts with verification constraints.
- Selective Verification reduces redundant target forward computations, contributing to the final speedup.
- Using 10 layers for the draft model was found to offer the best trade-off between acceptance length and decoding overhead.

Significance

The paper claims that FlexDraft offers a practical path toward efficient LLM inference under diverse serving conditions. By combining lightweight parameter adaptation with dynamic execution strategies, it successfully addresses the inherent trade-offs of parallel speculative decoding. The framework achieves high acceleration without compromising the target model's distribution (lossless) and avoids the quality degradation often associated with methods requiring extensive continual pretraining. The ability to maintain high throughput across varying batch sizes makes it particularly suitable for production environments where workload characteristics fluctuate.