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

Abstract: Block-diffusion drafters have recently emerged as a powerful alternative for speculative decoding by predicting multiple future-token distributions in a single parallel step. However, since these parallel predictions are sampled from position-wise marginals rather than fully conditioned sequences, committing to a single greedy path often fails to capture the target model's preferred trajectory. To address this, we propose BASTION, a budget-aware speculative decoding framework with tree-based diffusion drafting. Unlike existing methods that rely on static tree topologies, BASTION dynamically constructs query-dependent trees by balancing draft quality against hardware constraints. Our framework integrates three synergistic components: (1) an acceptance surrogate that estimates expected accepted length via path confidence, (2) an online latency estimator that calibrates a hardware-aware roofline model, and (3) an adaptive best-first expansion that grows the tree until marginal gains no longer justify incremental verification costs. BASTION is training-free, preserves the target model's distribution, and requires no per-setting tuning. Across diverse benchmarks and GPU architectures, BASTION achieves up to a 6.61x speedup over standard autoregressive decoding, outperforming state-of-the-art block-diffusion baselines by 39%.

• The paper introduces Bastion, a framework that uses budget-aware, adaptive tree-structured block diffusion drafting to accelerate language model inference.
• It employs a best-first, hardware-aware tree expansion strategy guided by path-confidence surrogates to optimize draft acceptance and verification costs.
• Empirical results show up to 6.61× speedup over autoregressive decoding, outperforming strong baselines across various GPU platforms and applications.

Bastion: Adaptive Speculative Decoding via Tree-Based Block Diffusion Drafting

Overview and Motivation

Autoregressive (AR) decoding remains a significant bottleneck for LLM inference, as the generation of each token mandates a full forward pass through the target model. Recent advances in speculative decoding, particularly with block diffusion drafters, attempt to alleviate this latency by proposing multiple token candidates in a parallel step, subsequently verified by the target model. However, conventional speculative decoding approaches often restrict themselves to sequential or greedy path extraction from the draft model's parallel outputs, underutilizing the drafter's multi-token distributions and failing to capture the target model's preferred continuations.

The Bastion framework introduces a budget-aware, training-free speculative decoding methodology that leverages an adaptive tree-structured block-diffusion drafting mechanism. Unlike previous solutions with fixed or static tree topologies, Bastion dynamically constructs query-dependent candidate trees, optimizing the trade-off between draft quality and device-constrained verification costs. This is achieved through a synergy of a drafter alignment-based acceptance surrogate, hardware-calibrated latency prediction, and a best-first tree expansion policy that halts when further tree growth yields diminishing expected speedup.

Algorithmic Components and System-Level Design

The Bastion framework is built upon three interacting modules:

1. Acceptance Surrogate from Drafter Logits: Bastion estimates the expected number of tokens likely to be accepted by the target model (i.e., accepted path length) through path-confidence scores derived solely from drafter marginals. Given the parallel marginal distributions of a block-diffusion drafter, every root-to-leaf path through the constructed prefix tree is scored as the cumulative product of token probabilities (the "path score"). The expected committed length is then the sum of these path scores over the tree, exploiting the position-wise independence induced by the block diffusive draft.
2. Best-First, Hardware-Aware Tree Expansion:

Traditional beam search or fixed-topology expansion spreads the verification budget either depth-wise or breadth-wise uniformly. Bastion instead employs a globally best-first expansion pattern, always adding the available node with maximum path score, thereby concentrating verification resources on the most probable, high-value continuations. This expansion is prefix-closed to ensure that candidate paths can always be found for continuity (Figure 1).

Figure 1: Drafter logits are expanded into an adaptive tree by best-first selection, directly optimizing estimated speedup given hardware constraints and draft likelihood.

1. Online and Analytically Calibrated Latency Model:

To avoid overcommitting compute to verification, Bastion incorporates an adaptive cycle-level controller. This controller uses a hardware-aware roofline model—calibrated online or offline to observed runtimes—to analytically predict the cost of verifying a tree of given size and structure. As the tree expands, the controller evaluates the estimated speedup, halting expansion at the unimodal optimum (i.e., when acceptance gain per cost is maximized).

Figure 2: Empirical acceptance-versus-latency curves exhibit saturation in committed length but superlinear latency growth with tree size, motivating the need for adaptive budget control.

Empirical Results and Ablations

Bastion delivers robust, high-magnitude speedups across multiple model-drafter pairs, benchmarks, and GPU classes. On Qwen3-8B with an A100 GPU, Bastion achieves a 6.61× average end-to-end speedup over AR decoding and consistently dominates strong baselines such as DFlash and EAGLE-3. This performance corresponds to a 39\% improvement over DFlash in comparable configurations.

Figure 3: End-to-end decoding speedup: Bastion consistently outperforms strong block-diffusion (DFlash) and tree-based (EAGLE-3) speculative decoding baselines across math, code, and chat workloads.

Results generalize to other hardware targets (H100, RTX PRO 6000 Blackwell, A6000) and extend across diverse domains, including mathematical reasoning, code generation, and instruction following. The empirical gains remain robust to decoding temperature, with only minor degradation under stochastic sampling.

Figure 4: Bastion achieves state-of-the-art wall-clock speedups across four major NVIDIA GPU platforms and three target models, outperforming all baselines per (model, GPU) pair.

Ablative experiments emphasize that best-first tree expansion substantially outperforms both greedy no-tree baselines and fixed beam search, even under matched verification budgets. The higher average acceptance lengths confirm the efficacy of path-confidence-guided node allocation.

Figure 5: Under fixed budgets, best-first expansion yields higher speedups and longer accepted continuations compared to rigid beam search or greedy no-tree baselines.

Additionally, budget sweeps and latency ablations confirm the fragility of static tree-size tuning. Bastion's online adaptive controller tracks or matches per-benchmark oracle settings, thereby substantially reducing the operational and deployment complexity—important for real-world inference systems.

Figure 6: Bastion’s adaptive budget selection nearly matches per-benchmark oracles and avoids the sharp speedup dropoff from over- or under-sizing the verification tree.

Verification latency modeling is supported by a lightweight roofline-based estimator, with linear calibration reducing root mean-squared error by $87-92\%$. Online variants maintain robust performance when offline calibration is unavailable.

Figure 7: Calibration and online adaptation align latency predictions with hardware-observed verification times, ensuring accurate real-time speedup estimation and robust controller behavior.

Technical Implications and Theoretical Properties

Bastion illustrates several technical advances for the speculative decoding ecosystem:

• Path-Score Surrogate:

The acceptance surrogate, derived from drafter-side self-verification, is shown to tightly correlate with target-model acceptance empirically, underpinning the utility of drafter marginals for speculative rollouts without additional alignment training or fine-tuning.

• Optimality of Best-First Expansion:

Theoretical results guarantee that best-first prefix-closed tree construction is provably optimal under the path-confidence acceptance surrogate for any fixed verification budget, exploiting the monotonicity and concavity of acceptance gain.

• Hardware Adaptivity and System Robustness:

The convexity of the latency-cost function, paired with the concavity of acceptance curves, ensures the unimodality of estimated speedup. This property allows for an efficient online search over tree budgets.

• Generalization Across Models and Hardware:

By decoupling tree construction from model- or platform-specific details and leveraging system-aware measurement, Bastion is robustly deployable across LLM and hardware variations.

• No Training or Tuning Required:

Bastion is fully stateless at deployment—no per-model training, no hyperparameter tuning, and no reliance on additional learning loops. The entire system is driven by deterministic expansion and runtime observation.

Practical and Theoretical Implications; Future Directions

Bastion's framework reinforces the importance of holistic, hardware-aware speculative decoding systems, where optimal utilization of parallel compute is mediated dynamically by real-time measurements and systematic surrogate modeling. The approach opens several avenues for future research:

• Integration with Advanced Multi-Token Drafters:

Extending Bastion's tree construction to other multi-token prediction paradigms (e.g., masked language modeling, hybrid AR-diffusion) could yield further system-level gains and generality.

• Finer-Grained Budgeting Policies:

Incorporating end-to-end client workload patterns, batch-level interactions, or multi-user scheduling may further enhance throughput and latency profiles in diverse serving regimes.

• Adaptive Alignment Mechanisms:

Though Bastion sidesteps the need for explicit target-drafter alignment, further improvements in acceptance surrogates or learning-to-align draft and target distributions online could improve the match in less aligned regimes.

• Optimized Verification Attention Kernels:

Systems engineering targeting tree-verification bottlenecks (e.g., custom CUDA attention masks, kernel fusion) remains a promising direction to close the gap between analytical and observed latency.

• Theoretical Extensions:

Exploring surrogate-driven expansion policies under distributional shift or adversarial drafter-target mismatches would enrich the theoretical understanding of speculative decoding optimality and robustness.

Conclusion

Bastion advances the speculative decoding literature by introducing a budget-aware, adaptive tree-based block-diffusion drafting mechanism that dynamically optimizes speculative rollouts under hardware constraints. Through provable optimality in tree selection, carefully crafted acceptance surrogates, and empirically calibrated latency models, Bastion achieves up to 6.61× speedup over autoregressive decoding and outperforms strong baselines without additional training or tuning. The architectural generality and hardware adaptivity presented here are likely to inform future LLM deployment strategies and prompt further investigation into system-aware, learning-agnostic inference acceleration techniques.