Bastion: Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting
This work addresses the problem of inefficient speculative decoding in large language model inference, offering a training-free method that improves speed while preserving output distribution.
BASTION introduces a budget-aware speculative decoding framework with tree-based diffusion drafting that dynamically constructs query-dependent trees to balance draft quality and hardware constraints, achieving up to 6.61x speedup over autoregressive decoding and outperforming state-of-the-art block-diffusion baselines by 39%.
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%.