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

@misc{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},
  howpublished = {Preprint},
  year         = {2026},
  url          = {https://arxiv.org/abs/2605.29727}
}