FlowSpec: Continuous Pipelined Speculative Decoding for Efficient Distributed LLM Inference

TL;DR

FlowSpec targets efficient edge deployment of autoregressive LLM inference under sparse requests by combining pipeline-parallelism with speculative decoding. It introduces three mechanisms—score-based step-wise draft verification, real-time draft pruning, and dynamic draft expansion—to maintain high pipeline utilization and forward progress in the draft tree. Empirical results on a real edge testbed show FlowSpec achieving up to $1.73\times$ speedups over strong baselines across multiple models and tasks, with robustness to bandwidth variations. The work offers a practical pathway to real-time, privacy-preserving edge inference and highlights areas for future integration with other model-parallelism techniques to scale further.

Abstract

Distributed inference serves as a promising approach to enabling the inference of large language models (LLMs) at the network edge. It distributes the inference process to multiple devices to ensure that the LLMs can fit into the device memory. Recent pipeline-based approaches have the potential to parallelize communication and computation, which helps reduce inference latency. However, the benefit diminishes when the inference request at the network edge is sparse, where pipeline is typically at low utilization. To enable efficient distributed LLM inference at the edge, we propose \textbf{FlowSpec}, a pipeline-parallel tree-based speculative decoding framework. FlowSpec incorporates three key mechanisms to improve decoding efficiency: 1) score-based step-wise verification prioritizes more important draft tokens to bring earlier accepted tokens; 2) efficient draft management to prune invalid tokens while maintaining correct causal relationship during verification; 3) dynamic draft expansion strategies to supply high-quality speculative inputs. These techniques work in concert to enhance both pipeline utilization and speculative efficiency. We evaluate FlowSpec on a real-world testbed with other baselines. Experimental results demonstrate that our proposed framework significantly improves inference speed across diverse models and configurations, achieving speedup ratios 1.37$\times$-1.73$\times$ compared to baselines. Our code is publicly available at \href{https://github.com/Leosang-lx/FlowSpec#}{https://github.com/Leosang-lx/FlowSpec\#}.

Figures (6)

• Figure 1: Left: Prefilling and decoding of LLM inference. Middle: Tensor Parallelism (TP). Right: Pipeline Parallelism with multi-/single-request. L1-L4 here represent the hidden layers.
• Figure 2: Overview of FlowSpec: An SD round includes a draft initialization step and multiple continuous SD steps. In draft initialization, segments of the initial generated draft tree are fed to V1 to fill all pipeline stages. Once the first subsequence completes all inference stages, the subsequent continuous SD steps begin. In each step, the current tree is pruned based on the accepted tokens, and then expanded to generate new draft tokens that serve as input for the next step.
• Figure 3: An example of tree pruning showing by the indices of the draft tokens. The above part is of the tree form, and the below is the sequence. In this case, $\mathcal{I}_\text{acc}=\{0,1\}$ and $\mathcal{I}_\text{draft}=\{3,6,7\}$. To prune the attention mask, only the entries whose column index and row index match the indices in $\mathcal{I}_\text{draft}$ are retained to form the updated mask.
• Figure 4: An example explaining merging the old pruned draft tree and the newly generated draft tree.
• Figure 5: FlowSpec's $\xi$ under different bandwidths for LLaMA2-chat on MT-bench.