SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification

Xupeng Miao; Gabriele Oliaro; Zhihao Zhang; Xinhao Cheng; Zeyu Wang; Zhengxin Zhang; Rae Ying Yee Wong; Alan Zhu; Lijie Yang; Xiaoxiang Shi; Chunan Shi; Zhuoming Chen; Daiyaan Arfeen; Reyna Abhyankar; Zhihao Jia

Source-linked AI summary

SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification

Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, Chunan Shi, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, Zhihao Jia

arXiv:2305.09781v4May 16, 2023cs.CLcs.DCcs.LG

TL;DR

Serving large language models is costly and slow because their size and autoregressive decoding constrain computational efficiency. SpecInfer uses small speculative models to build diverse token trees and verifies them in parallel with the LLM, achieving 1.5-2.8× gains for distributed inference and 2.6-3.5× gains for offloading-based inference while preserving generative accuracy.

Problem
Large language models are difficult to serve quickly and cheaply because their large parameter counts and autoregressive decoding limit serving efficiency.
Method
SpecInfer uses small speculative models to construct diverse token trees and verifies all candidate sequences in parallel with a tree-based decoding mechanism.
Results
1.5-2.8× for distributed inference and 2.6-3.5× for offloading-based inference: SpecInfer outperforms existing systems while preserving the same generative accuracy.
Takeaways & Limitations
Tree-based speculative inference and verification reduces LLM parameter accesses and end-to-end inference latency for distributed and offloading-based serving.
Takeaways & Limitations
SpecInfer’s performance improvement decreases as batch size increases, so it is most beneficial for low-latency inference.

Abstract

from arXiv · show

This paper introduces SpecInfer, a system that accelerates generative large language model (LLM) serving with tree-based speculative inference and verification. The key idea behind SpecInfer is leveraging small speculative models to predict the LLM's outputs; the predictions are organized as a token tree, whose nodes each represent a candidate token sequence. The correctness of all candidate token sequences represented by a token tree is verified against the LLM in parallel using a novel tree-based parallel decoding mechanism. SpecInfer uses an LLM as a token tree verifier instead of an incremental decoder, which significantly reduces the end-to-end latency and computational requirement for serving generative LLMs while provably preserving model quality. Our evaluation shows that SpecInfer outperforms existing LLM serving systems by 1.5-2.8x for distributed LLM inference and by 2.6-3.5x for offloading-based LLM inference, while preserving the same generative performance. SpecInfer is publicly available at https://github.com/flexflow/FlexFlow/

1 Introduction

SpecInfer addresses the latency and efficiency limits of autoregressive LLM serving by organizing diverse speculative predictions into token trees and verifying them in parallel. It preserves generative accuracy while accelerating distributed and offloading-based inference.

Serving challenges: Existing incremental decoding limits within-request parallelism and GPU utilization because each generated token depends on previously generated tokens.It also requires repeated access to the LLM’s parameters and past attention keys and values.
Prior approach: Sequence-based speculative inference uses one small speculative model to propose tokens, but model-capacity gaps can reduce alignment with the LLM.SSMs are generally orders of magnitude smaller than the LLM to reduce memory and runtime overheads.
SpecInfer approach: 52-57% to 96-97%: tree structures improve the success rate of verifying a token for stochastic decoding compared with sequence-based speculation.All candidate token sequences are verified against the LLM in parallel.
SpecInfer approach: SpecInfer organizes diverse speculation candidates from small speculative models into a token tree rather than considering only one sequence.Expansion-based and merge-based construction exploit diversity within one SSM and across multiple SSMs.
Correctness: SpecInfer’s multi-step speculative sampling guarantees equivalence with incremental decoding while maximizing the number of speculated tokens that can be verified.The mechanism targets stochastic decoding, where the next token is sampled from a probability distribution.
Results: 1.5-2.8× for distributed inference and 2.6-3.5× for offloading-based inference: SpecInfer outperforms existing systems while preserving the same generative accuracy.The system uses tree-based speculative inference and verification for both settings.

2 SpecInfer’s Overview

SpecInfer uses a learning-based speculator to construct token trees from speculative models, then uses the LLM as a parallel verifier rather than an incremental decoder. This reduces parameter accesses and latency while preserving the token dependencies needed for verification.

Speculative inference: The learning-based speculator takes a token sequence as input and produces a speculated token tree intended to overlap with the LLM’s incremental-decoding output.The token tree represents multiple candidate continuations for the input prompt.
Speculative models: SpecInfer can use pre-trained smaller architectures from the same model family as small speculative models, such as OPT-125M and OPT-350M for OPT-175B.These smaller models can be directly used as SSMs when trained on the same datasets.
Token-tree construction: Expansion-based and merge-based mechanisms construct token trees by exploiting diversity within a single SSM and across multiple SSMs.The overview contrasts these tree-construction mechanisms with the learning-based speculator and token-tree verification.
Tree verification: Existing serving systems incrementally decode one token at a time, whereas SpecInfer verifies all tokens in a token tree in one LLM decoding step.Each token’s activations account for all of its ancestors as preceding tokens.
Advantages: Reduced memory accesses to LLM parameters and reduced end-to-end inference latency are the two stated advantages over incremental decoding.These benefits target both distributed inference and offloading-based inference.

3 Learning-based Speculator

SpecInfer’s learning-based speculator constructs diverse token-tree candidates from one or more small speculative models, then verifies them against the LLM in parallel. It uses expansion- and merge-based construction strategies while leaving dynamic expansion for future work.

Single-sequence speculation can lose alignment exponentially with expected length, while verifying only one candidate token per step limits coverage.
SpecInfer aggregates predictions from one or more speculative models and organizes their diverse outputs into a token tree.The token tree represents candidate token sequences through nodes and supports expansion- and merge-based construction.
Expansion-based token tree construction: Expansion-based construction selects multiple top-k tokens from a speculative model, exploiting the observation that the LLM token often appears among these candidates.
Dynamic expansion of token trees remains an open research problem outside this paper’s scope.
Merge-based token tree construction: Merge-based construction combines multiple speculative-model outputs into a tree containing their token sequences.SpecInfer adaptively boost-tunes multiple speculative models using general text data, while parallel execution avoids increasing speculative latency or per-GPU memory overhead.

4 Token Tree Verifier

SpecInfer verifies all candidate sequences in a speculated token tree against the LLM in parallel, using tree attention and tree-based decoding to reduce redundant computation and memory accesses. It supports greedy and stochastic verification, with multi-step speculative sampling preserving incremental decoding’s generative performance.

Token Tree Verification: SpecInfer’s token tree verifier checks all candidate sequences in parallel with one pass over the LLM parameters.This opportunistically decodes multiple tokens per step instead of using incremental decoding for one token at a time.
Tree Attention: Tree attention generalizes sequence attention by computing each node’s output from the token sequence represented by that node.For a merged token tree, this produces attention outputs for all covered token sequences.
Tree-Based Parallel Decoding: Shared-prefix sequences make separate key-value caches redundant and add kernel-launch overhead in straightforward sequence-based verification.Causal masking gives shared-prefix tokens identical attention outputs, motivating shared cache reuse and fused computation.
Tree-Based Parallel Decoding: SpecInfer reuses one key-value cache across token sequences by traversing the token tree depth-first.The traversal updates the shared cache while preserving the preceding-token keys and values needed for each new token.
Tree-Based Parallel Decoding: Topology-aware causal masking fuses attention computation for all tree tokens into a single kernel.The tree topology stores verified and speculated tokens together in the key-value cache for batched attention computation.
Token Verification: Multi-step speculative sampling verifies stochastic token trees while provably preserving the LLM’s generative performance as incremental decoding.The method performs multi-step verification across branches and supports both merge-based and expansion-based speculation.

5 System Design and Implementation

SpecInfer combines small speculative models, token-tree construction, and parallel LLM verification within a distributed runtime. Its design reorganizes inference to exploit parallelism while keeping request-management and speculation overheads comparatively small.

Runtime workflow: SpecInfer iteratively schedules requests, gathers SSM predictions, builds token trees, and verifies them against the LLM.The request manager adapts iteration-level scheduling, while SSMs speculate and the LLM performs tree-based parallel decoding.
Parallelization: SSMs use data parallelism, while the LLM combines tensor model parallelism within nodes with pipeline model parallelism across nodes.This hybrid strategy distributes Transformer layers across GPUs and nodes for LLM serving.
Runtime overhead: Request scheduling, token-tree merging, verification, and token-only communication add negligible overhead compared with LLM inference execution.The runtime avoids transferring token vector representations between the request manager and GPU workers.
Scheduling: Continuous batching schedules LLM execution by iteration, allowing newly arrived requests to begin without waiting for all current requests to finish.After each iteration, SpecInfer checks request status and returns completed results.
Overheads: SpecInfer’s additional memory and computation overheads are generally one or two orders of magnitude smaller than executing LLM inference.Memory overhead includes SSM parameters and token-tree verification, while computation overhead includes SSM generation and verification of unmatched candidates.
Target applications: SpecInfer targets distributed and offloading-based LLM inference, where memory capacity and CPU–GPU data movement constrain serving.Offloading can benefit because verifying multiple tokens reduces decoding steps and communication between CPU DRAM and GPU HBM.

6 Evaluation

The evaluation measures SpecInfer across LLaMA and OPT models, multiple datasets, parallel configurations, and decoding settings. It reports substantial latency reductions while preserving generated sequences, with benefits depending on batch size and deployment scenario.

Distributed LLM inference: SpecInfer outperforms incremental decoding by 1.5-2.5× for single-node multi-GPU inference and 2.4-2.8× for multinode multi-GPU inference.It generates the exact same token sequences as incremental decoding for all prompts.
Speculation comparison: Compared with sequence-based speculative inference, tree-based SpecInfer further reduces LLM serving latency by 1.2-1.5×.The paper attributes this to token-tree optimization, parallel verification, and multi-step speculative sampling.
Evaluation caveat: SpecInfer’s performance advantage decreases as batch size increases, making it most beneficial for low-latency inference.Larger batches provide incremental decoding with more parallelizable computation and leave fewer spare GPU resources for tree-based decoding.
Offloading-based LLM inference: SpecInfer reduces per-token latency versus FlexGen by 2.6-3.5× for offloading-based inference on OPT-13B and OPT-30B using one 24GB A10 GPU.The improvement is associated with fewer LLM decoding steps and CPU–GPU parameter transfers.
Token-tree analysis: Token trees reduce LLM decoding steps by 1.2-1.5× for greedy decoding and 1.3-1.4× for stochastic decoding.The comparison uses sequence-based speculation as the tree-width = 1 baseline.
Token-tree analysis: Larger token-tree widths reduce latency for batch sizes BS = 1 and 2 but increase verification latency for BS ≥4.The trade-off reflects available sparse GPU resources at small batch sizes versus increased tree-verification work at larger batches.

7 Related Work

Related work spans lossless and lossy acceleration, hardware-level ML optimization, tree-structured attention, and multi-sample decoding. SpecInfer differs by organizing multiple speculative candidates for parallel verification without directly reducing LLM computation.

Lossless and lossy acceleration: Lossless acceleration methods use an LLM as a verifier, while sequence-based speculative inference checks tokens proposed by a small speculative model.SpecInfer extends this direction with tree-based candidate organization and verification.
Hardware optimization: SpecInfer’s techniques are orthogonal to automatic kernel generation and graph-level DNN transformations and can be combined with them.The paper identifies such combinations as a future direction.
Lossless and lossy acceleration: Lossy acceleration methods, including BiLD and model compression, trade possible token-quality degradation or predictive performance for lower latency or resource requirements.Compression approaches include quantization and structured pruning.
SpecInfer’s distinction: SpecInfer reorganizes LLM inference into a more parallelizable form rather than directly reducing computation, accepting manageable memory and computation overheads.This distinction separates it from compression-based acceleration.
Tree-structured attention: Unlike tree-structured attention, SpecInfer uses trees to combine candidate sequences and decode speculative alternatives against an LLM.The related technique instead represents a single input text’s hierarchical parse structure.
Multi-sample decoding: Unlike beam search and top-k/top-p sampling, SpecInfer uses SSMs to speculatively predict candidates and an LLM to verify them in parallel.Its objective is reducing decoding iterations and latency through speculative verification.

8 Conclusion

SpecInfer accelerates generative LLM inference by organizing diverse SSM predictions into token trees and verifying them in parallel. The approach reduces memory accesses and end-to-end latency in distributed and offloading-based serving.

Conclusion: SpecInfer organizes diverse speculation candidates into token trees and verifies them against the LLM with tree-based parallel decoding.The conclusion frames this as the system’s central insight for predicting LLM outputs.
Conclusion: SpecInfer reduces LLM parameter memory accesses and end-to-end inference latency for distributed and offloading-based LLM inference.The stated scope covers both deployment scenarios.

A.1 Abstract

The artifact includes the SpecInfer implementation and the materials needed to reproduce the paper’s experiments.

The artifact contains code for running SpecInfer.
It includes datasets for reproducing the experiments.
It provides scripts to reproduce the paper’s experiments.

A.2 Artifact check-list (meta-information)

The checklist specifies the artifact’s algorithm, programs, build and runtime requirements, hardware, metric, outputs, and evaluation settings.

The artifact implements tree-based speculative inference.
The listed programs are spec_infer.cc and incr_decoding.cc, compiled with CMake.
The runtime uses CUDA, NCCL, MPI, UCX, and Python3 on two AWS g5.12xlarge instances.
Evaluation measures end-to-end average latency for server-grade GPU and offloading-based inference.

A.3 Description

The artifact provides installation instructions, supported models, experiment scripts, and output-processing tools for reproducing SpecInfer evaluations. The workflow covers server-grade GPU and offloading-based inference, with configurable hardware, models, datasets, precision, batching, and generation length.

Description: The repository contains SpecInfer’s source code, build instructions, and scripts for reproducing the experiments.
Installation: The recommended setup uses two AWS g5.12xlarge instances, each with four NVIDIA A10 GPUs, 48 CPU cores, and 192 GB DRAM.
Installation: The software prerequisites include CUDA 12.1, NCCL, Rust, CMake, Python3, and additionally UCX and MPI for multinode experiments.
Experiments: The experiments use LLaMA-7B, OPT-30B, and LLaMA-65B as LLMs, with LLaMA-68M and OPT-125M as speculative models.
Experiments: Server-grade GPU evaluation compares incremental, sequence-based speculative, and tree-based speculative decoding across one- and two-node configurations.
Results and outputs: The scripts generate raw outputs and logs, convert results to CSV for plotting, and should produce performance similar to Figures 7–11, with expected variability.
Configuration: Evaluation configurations can change GPU and CPU resources, memory, batch size, models, prompt dataset, precision, and maximum generated tokens.