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.61× speedup over standard autoregressive decoding, outperforming state-of-the-art block-diffusion baselines by 39%.

BibTeX

@inproceedings{oh2026bastion,
  title     = {BASTION: Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting},
  author    = {Oh, Soowon and Cao, Nam and Kim, Yujin and Jung, Hojung and Ahmad, Huzama and Bae, Sangmin and Yun, Se-Young},
  booktitle = {AdaptFM @ ICML},
  year      = {2026},
  url       = {https://arxiv.org/abs/2605.29727},
  note      = {Oral Presentation}
}