🤖 machine learning

Bastion: Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting

BASTION is a training-free, budget-aware speculative decoding framework that dynamically constructs query-dependent tree structures via adaptive best-first expansion to balance draft quality with hardware constraints, achieving up to a 6.61x speedup over standard autoregressive decoding while outperforming existing block-diffusion baselines.

Original authors: Soowon Oh, Nam Cao, Yujin Kim, Hojung Jung, Huzama Ahmad, Sangmin Bae, Se-Young Yun

Published 2026-05-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Soowon Oh, Nam Cao, Yujin Kim, Hojung Jung, Huzama Ahmad, Sangmin Bae, Se-Young Yun

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 story with a very wise, but incredibly slow, editor (the Target Model). Every time you write a single word, you have to wait for the editor to read it, think about it, and give you the green light before you can write the next one. This is how current AI chatbots work: they write one word at a time, waiting for a "check" after every single step. It's accurate, but it's painfully slow.

To speed this up, researchers use a Drafting Model (a faster, less wise assistant) to guess the next few words before the editor checks them. If the editor agrees with the guess, you get to write multiple words at once, skipping the waiting time.

However, there's a catch with the newest, fastest drafters (called Block Diffusion). Instead of guessing a sentence one word at a time, they shout out a whole block of words all at once. The problem is, because they shout them all together, they aren't always sure how the words fit together in a logical sequence. It's like a chef throwing a bunch of ingredients onto a counter at once; individually, they look good, but if you just grab the top one, you might end up with a weird, nonsensical dish.

BASTION is a new system designed to fix this mess and make the whole process lightning fast. Here is how it works, using simple analogies:

1. The "Tree of Possibilities" (Instead of a Single Path)

Old methods would take the drafter's shout-out, pick the single "best" word, then the next "best" word, and hope for the best. If that path turned out to be wrong, the editor had to reject the whole thing, and you wasted time.

BASTION is different. Imagine the drafter doesn't just give you one path; it gives you a family tree of possibilities.

At the first step, it says, "Maybe the next word is 'cat' (80% chance) or 'dog' (20% chance)."
Instead of just picking 'cat', BASTION builds a tiny tree: one branch for 'cat', one for 'dog'.
Then, for the next word, it branches out again from both 'cat' and 'dog'.
Suddenly, you have a small forest of potential sentences growing from the same starting point.

2. The "Smart Gardener" (Budget-Aware Controller)

Here is the tricky part: You can't grow an infinite forest. The editor (Target Model) has a limit on how many branches it can check at once. If you grow a tree that is too wide or too deep, the time it takes to check the tree becomes longer than just writing the words one by one.

This is where BASTION's "Smart Gardener" comes in.

The Budget: The Gardener knows exactly how much time the editor has (the "budget").
The Strategy: Instead of growing the tree to a fixed size (like "always 10 branches"), the Gardener looks at the confidence of each branch.
- If a branch looks very promising (high confidence), the Gardener grows it deeper.
- If a branch looks weak, it stops growing there.
The Stop Sign: The Gardener constantly asks, "Is adding one more branch going to give us more speed, or is it just going to waste time checking dead ends?" As soon as the cost of checking a new branch outweighs the benefit, the Gardener stops growing and sends the tree to the editor.

3. The "Speedometer" (Hardware Awareness)

Different computers (GPUs) are like different cars. A sports car (a powerful GPU) can check a huge tree very quickly. A compact car (a weaker GPU) might struggle with the same tree.

BASTION has a built-in Speedometer that knows exactly how fast your specific computer is. It doesn't just guess; it measures how long it takes to verify a tree of a certain size on your machine. It uses this real-time data to decide exactly how big the tree should be for your specific setup, ensuring you get the maximum speed without overloading your computer.

The Result

By combining these ideas, BASTION achieves what the paper calls a 6.61x speedup.

Standard AI: Writes 1 word, waits, writes 1 word, waits. (Speed: 1x)
Old Fast Methods: Guesses a few words, but often gets stuck on the wrong path. (Speed: ~2-3x)
BASTION: Grows a smart, custom-sized tree of guesses, checks the most promising paths, and stops exactly when it's most efficient. (Speed: ~6.6x)

In short, BASTION is like a smart project manager for AI writing. Instead of blindly guessing or building a rigid structure, it dynamically builds a flexible "tree of options" that is perfectly sized for the computer's speed, ensuring the AI writes as fast as possible without making mistakes.

Technical Summary: BASTION

Problem Statement

Large Language Models (LLMs) suffer from high serving costs due to the sequential nature of autoregressive (AR) decoding, which requires a separate target-model forward pass for every generated token. Speculative decoding addresses this by using a lightweight drafter to propose multiple candidate tokens, which the target model verifies in parallel.

While recent advances in block-diffusion drafters (e.g., DFlash, TiDAR) allow for the parallel prediction of multiple future tokens in a single step, they introduce a critical limitation: these models generate position-wise marginals rather than fully conditioned sequences. Consequently, committing to a single greedy (top-1) path often results in incoherent sequences under the target model's distribution. Existing tree-based speculative decoding methods (e.g., EAGLE-3) typically rely on static, heuristically predefined tree topologies that do not adapt to the specific query or hardware constraints, leading to suboptimal trade-offs between draft quality and verification overhead.

Methodology: BASTION

BASTION (Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting) is a training-free framework that dynamically constructs query-dependent tree topologies to maximize decoding speedup. It integrates three synergistic components:

1. Tree-based Block-Diffusion Drafting

Unlike autoregressive drafters that generate tokens sequentially, BASTION utilizes a block-diffusion drafter to produce position-wise marginal distributions $\{q_k(\cdot)\}$ in a single forward pass.

Topology Decoupling: The drafting cost ( $T_{draft}$ ) is invariant to the tree structure (chains, wide trees, or deep trees) because the block is processed in parallel.
Marginal Independence: Candidates at different positions are independent of each other given the validated context, allowing them to be combined freely into a prefix tree $\mathcal{T}$ .
Path Scoring: Each node $i$ in the tree at depth $d(i)$ is assigned a path score $\rho(i) = \prod_{k=1}^{d(i)} q_k(x_{ik})$ , representing the joint probability of the path under the drafter's distribution.

2. Acceptance Surrogate and Best-First Expansion

To determine which nodes to include in the verification tree without incurring the cost of target-model verification, BASTION employs a path-confidence surrogate.

Surrogate Estimation: The expected accepted length $\mathbb{E}[A(\mathcal{T})]$ is approximated by the sum of path scores: $\hat{A}(\mathcal{T}) = \sum_{i \in \mathcal{T}} \rho(i)$ . This relies on the premise that high-probability drafter paths align with the target model's preferred trajectory.
Optimal Construction: The paper proves that a best-first expansion strategy is provably optimal for maximizing $\hat{A}(\mathcal{T})$ under a fixed budget. The algorithm iteratively adds the valid node with the highest $\rho(i)$ , ensuring the tree grows along the most promising trajectories first. This yields a concave relationship between tree size and expected acceptance, exhibiting diminishing returns.

3. Hardware-Aware Online Budget Controller

BASTION dynamically determines the optimal tree size $N$ at each decoding step by balancing the marginal gain in acceptance length against the marginal increase in verification latency.

Latency Modeling: The framework uses a hardware-aware roofline model to estimate verification latency $T_{verify}(N)$ . This model accounts for memory bandwidth and compute constraints, calibrated via observed runtimes to correct for kernel fusion and scheduling overhead.
Adaptive Stopping: The controller expands the tree until the estimated speedup $S(N) = \frac{\hat{A}(\mathcal{T}_N) \cdot L_{AR}}{C(N)}$ begins to decrease. This ensures the system halts expansion exactly when the cost of verifying additional nodes outweighs the benefit of a longer accepted sequence.

Key Contributions

Tree-based Block Diffusion Drafting: Introduction of a best-first prefix-tree construction algorithm that leverages drafter-logit path scores to build high-quality, prefix-closed candidate trees from block-diffusion logits, moving beyond single-path greedy decoding.
Hardware-Aware Online Budget Controller: A novel controller that dynamically determines the optimal tree size at each step by balancing expected acceptance gains against hardware-calibrated verification latency, eliminating the need for static, fixed-size trees.
Extensive Empirical Validation: Comprehensive evaluation across diverse target models (Qwen3, Llama-3.1), benchmarks (math, code, chat), and GPU architectures (A100, H100, A6000, RTX PRO 6000), demonstrating significant speedups without per-setting tuning.

Experimental Results

Evaluated on an NVIDIA A100 GPU with greedy decoding (temperature 0):

Speedup: BASTION achieves an average 6.61× speedup over standard autoregressive decoding on Qwen3-8B.
Comparison to Baselines:
- Outperforms EAGLE-3 (state-of-the-art AR-based tree decoding) by 2.45× (6.61× vs. 2.70× average).
- Outperforms DFlash (state-of-the-art block-diffusion single-path) by 1.39× (6.61× vs. 4.76× average).
Acceptance Length: BASTION achieves significantly larger Average Accepted Lengths (AAL) compared to single-path baselines (e.g., 8.48 vs. 6.24 on Qwen3-8B), confirming that tree-structured verification successfully captures longer valid continuations.
Generalization: The performance gains are consistent across different GPU generations (including H100 and Blackwell architectures) and model families, indicating the method is not tied to specific hardware or model configurations.

Significance and Claims

The paper claims that BASTION represents a significant advancement in speculative decoding by resolving the tension between the parallel efficiency of block-diffusion drafters and the sequential coherence required by target models. By replacing static, heuristic tree shapes with a dynamic, budget-aware construction, BASTION fully utilizes available parallel compute resources while avoiding the latency penalties of over-verifying low-probability paths.

Crucially, the method is training-free and requires no per-setting tuning, making it immediately deployable across diverse hardware and model configurations. The authors assert that their approach effectively bridges the gap between the theoretical potential of block-diffusion drafting and practical, high-throughput inference, achieving near-optimal speedups by adaptively aligning draft structure with real-time system constraints.