SpecFold speeds up diffusion language model text generation by folding redundant work
SpecFold: Folding Multi-Branch Redundancy for Faster Speculative Decoding in Diffusion Language Models
Artificial IntelligenceComputation and LanguageMachine Learning
Summary
Diffusion language models create text step by step, which can be slow. The authors found that during the process, many parts of the work are very similar across different guesses the model makes at once. They designed SpecFold, a method that reuses these similar parts to speed things up without losing accuracy. This makes speculative text generation faster while keeping results just as good.
What this means in practice
- •For language model engineers: Speed up speculative decoding in diffusion language models by reusing computations across branches within a verification step using SpecFold.
- •For machine learning infrastructure teams: Enhance throughput of text generation systems running diffusion language models with SpecFold’s efficient multi-branch execution to reduce latency and compute costs.
Authors
Chung-En Ho, Weiyu Sun, Cheng-Jhih Shih, He Li, Yong Liu, Yingyan, Lin
Abstract
Diffusion large language models (DLLMs) generate text through iterative block denoising, and multi-branch speculative decoding accelerates this process by verifying a main branch together with multiple draft branches in a single forward pass. While prior DLLM acceleration methods primarily exploit temporal redundancy across denoising steps, we identify a complementary redundancy axis within each speculative verification step: multi-branch computational redundancy. During speculative verification, draft branches inherit most tokens from their parents while unmasking a small set of additional positions, causing large portions of hidden states to remain highly similar across branches. We propose SpecFold, an algorithm-system co-design that exploits this multi-branch redundancy to reduce the cost of multi-branch speculative verification. Algorithmically, SpecFold performs token-level residual gating and selectively reuses parent computation through folded attention and FFN while preserving residual hidden states. Systemically, a Triton kernel implementation translates this fine-grained reuse into end-to-end throughput gains through efficient sparse multi-branch execution. SpecFold is orthogonal to temporal caching and compatible with existing DLLM speculation strategies. Across two DLLM families, five models, and five standard benchmarks, SpecFold achieves up to 1.64x throughput over Spiffy and up to 1.99x over vanilla decoding, while maintaining comparable task performance.