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

In plain words

Writing text with a large language model usually means producing one token at a time, and each token needs a full pass over the model's weights. That is slow, and it is slower still when the weights sit in ordinary memory and have to be copied to the GPU for every step.

SpecInfer speeds this up by letting smaller, cheaper models guess ahead. Their guesses form a tree of possible continuations rather than a single line. The big model then checks every branch of the tree in one pass, and the tokens that agree with what it would have produced anyway are kept. The output matches ordinary decoding, so the speed-up is not paid for in output quality.

The authors report that SpecInfer outperforms existing serving systems by 1.5 to 2.8 times when the model is spread over several GPUs, and by 2.6 to 3.5 times when the model is offloaded to CPU memory. Those figures are their own measurements.

TL;DR

SpecInfer replaces one-token-per-pass decoding with speculate-then-verify. Small models build a token tree, and the LLM scores all of its nodes in a single pass using tree attention. Greedy verification reproduces incremental greedy decoding, and the paper adds a multi-step sampling rule with a stated correctness theorem. The author-reported gains are 1.5-2.8x for distributed serving and 2.6-3.5x for offloading. The released inference branch implements only greedy verification and one draft model.

Problem

• Incremental decoding runs one LLM decoding step per generated token, and each step reads all model parameters §2 alg. 1.
• Tokens depend on earlier tokens, so the usual serving loop cannot produce several tokens of one request in parallel §2.
• With offloading, parameters are moved from CPU DRAM to the GPU for each step, so transfers dominate the cost §2 §6.3.
• Cached keys and values take memory and limit how many requests fit in a batch §1.
• Existing speculative inference uses one small speculative model (SSM) to propose one token sequence per step. The chance that a whole sequence matches decays with its length, and a single SSM is bounded by the capacity gap to the LLM §1 §3.
• Speculation helps only when its tokens overlap the LLM's actual output. If the overlap is empty, the step gains nothing §2.

Method

Overview. A learning-based speculator takes the current token sequence and produces a token tree. The LLM is then run once over the whole tree (tree-parallel decoding), and a verifier accepts a prefix of the tree §2 alg. 2 fig. 2. The LLM is used as a verifier, not as a next-token decoder §2.

Speculator (token trees). A token tree is a tree whose nodes carry tokens; each node stands for the token sequence on its path from the root §3.

1. Expansion-based construction: one SSM expands the tree. A configuration <k_1, ..., k_m> gives the number of children to expand at each of up to m steps. The example <2, 2, 1> yields four sequences §3 fig. 3. Expansion is static; dynamic expansion is left open §3.
2. Acceptance by top-k: if the SSM's top-k tokens contain the LLM's choice, the token counts as matched. For LLaMA-7B with LLaMA-68M, top-5 raises the match rate from 70% to 89% (greedy) and from 57% to 97% (stochastic) §3 tab. 1.
3. Merge-based construction: several SSMs each give a sequence or tree, and the token tree merge takes the union of their sequences §3, Definition of token tree merge.
4. Collective boost-tuning: SSMs are fine-tuned one at a time on a text corpus (OpenWebText is the example). Prompts where an SSM already matches the LLM are removed before the next SSM is tuned, using adaptive boosting §3. Details are referred to the authors' earlier work §3 §6.1.

Tree attention. For node u, tree attention is ordinary attention computed over the token sequence S_u of that node §4.1. The causal mask is

\mathrm{mask}(A)_{jk} = \begin{cases} A_{jk} & j \geq k \\ -\infty & j < k \end{cases} §4.1

Using -∞ rather than 0 makes the softmax weight of masked positions exactly 0 §4.1.

Tree-based parallel decoding. All nodes are computed in one LLM decoding step §4.2.

1. Shared key-value cache: a depth-first traversal of the tree writes keys and values into one cache, so different branches do not need separate caches §4.2 fig. 4.
2. Topology-aware causal mask: all tree tokens, plus the verified prefix token(s), go through one attention kernel. The mask is changed from the sequence order to the tree topology, so that a node sees its ancestors and itself, not its siblings' branches. The result matches incremental decoding exactly §4.2.

# Topology-aware causal mask  [§4.2]
# Input:  token tree N, number of verified prefix tokens P
# Output: boolean mask of shape (N, P + N)

def tree_mask(tree, prefix_len):
    # Every tree node sees all verified prefix tokens  [§4.2]
    mask = all_false(num_nodes(tree), prefix_len + num_nodes(tree))
    for node in tree:
        mask[node][:prefix_len] = True
        # A node sees itself and its ancestors, not other branches  [§4.2]
        for other in ancestors_and_self(tree, node):
            mask[node][prefix_len + other] = True
    return mask

Verification (Algorithm 2). The serving loop repeats speculate, verify, append alg. 2.

# SpecInfer serving loop  [alg. 2]
# Input:  prompt tokens I
# Output: generated tokens S, ending at EOS

def serve(prompt_tokens):
    # Start the current sequence S with the prompt  [alg. 2]
    sequence = prompt_tokens
    while True:
        # Small models expand a token tree N from S  [§2] [§3]
        tree = speculate(sequence)
        # One LLM pass gives one token O(u) per tree node  [§4.2] [alg. 2]
        outputs = tree_parallel_decode(llm, tree)
        # Pick the verified tokens V, greedily or by sampling  [§4.3]
        if use_greedy:
            verified = verify_greedy(outputs, tree)
        else:
            verified = verify_stochastic(outputs, tree)
        # Append the verified tokens; stop at end-of-sequence  [alg. 2]
        for token in verified:
            sequence.append(token)
            if token == EOS:
                return sequence

Greedy verification walks down the tree while a child's token equals the LLM's prediction at the current node §4.3 alg. 2.

# VerifyGreedy  [alg. 2] [§4.3]
# Input:  LLM outputs O (one token per node), token tree N
# Output: verified tokens V

def verify_greedy(outputs, tree):
    # Start at the root with nothing verified  [alg. 2]
    verified = []
    node = tree.root
    # Look for a child whose token is the LLM's prediction here  [§4.3]
    child = matching_child(tree, node, outputs[node])
    while child is not None:
        verified.append(tree.token[child])
        node = child
        child = matching_child(tree, node, outputs[node])
    # Append the LLM's own prediction at the last matched node  [alg. 2]
    verified.append(outputs[node])
    return verified

Stochastic verification is multi-step speculative sampling. Each child of the current node is tested in random order against the LLM distribution, using the distribution of the SSM that proposed it §4.3 alg. 2.

# VerifyStochastic  [alg. 2] [§4.3]
# Input:  LLM distributions per node, SSM distributions per node and SSM
# Output: verified tokens V

def verify_stochastic(llm_dist, ssm_dist, tree):
    # Start at the root with nothing verified  [alg. 2]
    verified = []
    node = tree.root
    # Descend while the current node has children  [alg. 2]
    while not is_leaf(tree, node):
        candidates = children(tree, node)
        accepted = False
        while candidates is not empty:
            # Pick one candidate at random; s is the SSM that proposed it
            child = pick_random(candidates)
            token = tree.token[child]
            s = tree.proposer[child]
            # Accept with probability min(1, p_llm / p_ssm_s)  [alg. 2]
            ratio = llm_dist[node][token] / ssm_dist[s][node][token]
            if random_uniform() <= ratio:
                verified.append(token)
                node = child
                accepted = True
                break
            # Reject: the LLM distribution at this node becomes the
            # normalised residual max(0, p_llm - p_ssm_s); drop s  [alg. 2]
            llm_dist[node] = normalise(
                max(0, llm_dist[node] - ssm_dist[s][node]))
            candidates.remove(child)
        if not accepted:
            break
    # Sample the next token from the LLM at the final node  [alg. 2]
    verified.append(sample(llm_dist[node]))
    return verified

The paper states that multi-step speculative sampling leaves the output distribution equal to that of the LLM (Theorem in §4.3). Its proof is not in the source; the source cites a separate reference that carries the same arXiv id §4.3 [reference: miao2023specinfer].

Runtime (§5).

• A request manager schedules at iteration granularity, adapting iteration-level scheduling from Orca §5.1.
• SSMs are small enough for one GPU each and run data-parallel across GPUs §5.1 fig. 6.
• The LLM uses tensor parallelism inside a node and pipeline parallelism across nodes, as in Megatron-LM §5.1.
• Continuous batching lets new requests start after any decoding iteration §5.1.
• The attention kernel is customised from FasterTransformer. Each thread block computes one head for one request §5.2.
• The paper estimates that one SSM adds less than 1% memory per GPU for SSMs 100 to 1000 times smaller than the LLM. It estimates that token-tree memory is small next to the key-value cache. These are estimates, not measurements §5.3.

Hyperparameters

Results

Author-reported. Setup: five prompt datasets, LLaMA-7B / OPT-13B / OPT-30B / LLaMA-65B as LLMs, LLaMA-68M and OPT-125M as SSMs, AWS g5.12xlarge with A10 GPUs, baselines vLLM, HuggingFace TGI, FasterTransformer (distributed) and FlexGen (offloading) §6.1 §6.2 §6.3.

Per-system bars for the distributed comparison appear only in fig. 7, not as numbers in the text.

Table 1: success rate of verifying a token with top-k tokens of LLaMA-68M, for LLaMA-7B tab. 1.

Table 2: average tokens verified per decoding step, LLaMA-7B with LLaMA-68M, speculation length 8, by tree width tab. 2.

Table 3: multi-step speculative sampling vs naive sampling, average verified tokens per stochastic step, width 5 and depth 8, LLaMA-7B with LLaMA-68M tab. 3.

Paper vs code

Searched: the abstract links github.com/flexflow/FlexFlow. The default branch at clone time is a later refactor (lib/) without SpecInfer code, so the clone was switched to the inference branch, which contains inference/spec_infer/ (code at 3a825ed, HEAD of that branch, later than the paper). The artifact repository named in the source's artifact appendix (goliaro/specinfer-ae) was not cloned. No collective boost-tuning code was found.

• Stochastic verification is absent. The paper's multi-step speculative sampling §4.3 alg. 2 is not in the released code. The only verification is greedy token matching: traverse_verify_tree walks the tree by exact match with the LLM's output request_manager.cc request_manager.cc. Searches for "stochastic" and "multi-step" in inference/, src/runtime/ and include/flexflow/request_manager.h found only do_sample in the incremental-decoding path incr_decoding.cc.
• Tree width cap. Tables 2 and 3 use widths up to 5 tab. 2 tab. 3. The code sets MAX_BEAM_WIDTH = 3 and asserts every pushed width is at most this value batch_config.h request_manager.cc. The released code cannot run the width-5 setting as written.
• Merge-based construction is disabled. merge_dfs_trees asserts a single SSM ("currently using one ssm") and returns that tree before the merge code runs request_manager.cc. The paper's merge-based construction is therefore not active in this code.
• Boost-tuning is not released. No boosting or SSM fine-tuning code was found in inference/, src/runtime/ or include/flexflow/ §3 §6.1.
• Expansion configuration matches. The paper's configuration <1, 1, 3, 1, 1, 1, 1, 1> §6.1 matches the driver: widths 1, 1 and expansion_degree (default 3) are pushed spec_infer.cc, and steps beyond the pushed list use width 1 request_manager.cc. Depth is capped at 8 batch_config.h. Passing -1 pushes no widths, so every step has width 1, which is sequence-based speculation.
• Tree size. The paper does not state a cap on tree size. The driver uses 23 spec_infer.cc. The batch structure allows up to 64 batch_config.h, and the bit-mask update asserts a tree of at most 64 tokens request_manager.cc.
• Causal mask. The tree attention kernel reads a per-request bit mask of tree ancestry and the non-tree cache size tree_inc_multihead_self_attention.cu, which matches the topology-aware mask described in §4.2.

Pitfalls

• Only greedy verification is released. An implementer who needs sampling must write multi-step speculative sampling from alg. 2 and §4.3; the released code does not provide it.
• Ambiguous residual sampling. The stochastic algorithm replaces the LLM distribution at a node by the residual after each rejection, then samples the final token from that distribution alg. 2. Read literally, the final token at a node where all children were rejected comes from the residual. The paper does not say this in words.
• Undefined edge cases in sampling. The pseudocode divides by the SSM probability and normalises the residual. What happens when the SSM probability is 0 or the residual sums to 0 is not stated.
• Width 5 is not reproducible from this code. See the width cap above. Tables 2 and 3 cannot be reproduced from the released inference branch as it stands.
• Multi-SSM gains are not evaluated here. All evaluation uses expansion-based construction with one SSM §6.1. Merge-based and boost-tuned results are deferred to the authors' earlier work §6.1.
• Greedy equivalence is by construction. The paper says the output is the same as incremental decoding for all prompts §6.2. Only the stochastic theorem §4.3 is a stated guarantee, and its proof is not in the source.
• Speedup depends on batch size. Gains shrink as the number of concurrent requests grows, because the spare GPU capacity that tree verification uses is smaller at larger batches §6.2. Wide trees hurt at large batch sizes; width 2 or 3 is best there §6.4.
• Over-generation is truncated. Verification can produce more than 128 tokens per iteration; the evaluation cuts the output to 128 §6.2.
• Abstract ranges mix settings. The abstract's 1.5-2.8x combines the single-node and multi-node ranges from §6.2. Use the section's ranges per setting.
• Baseline versions are not given. The text does not state the versions of vLLM, TGI, FasterTransformer or FlexGen used.
• Memory overhead is an estimate. The claims that one SSM adds under 1% memory and that tree memory is small next to the key-value cache are estimates §5.3, not measured in the text.
• Third-party weights. The models come from HuggingFace (LLaMA, OPT) §6.1; their licenses are separate from this paper's CC BY 4.0 license.

Lineage

• Builds on: sequence-based speculative decoding with one small model, Fast Inference from Transformers via Speculative Decoding and Accelerating Large Language Model Decoding with Speculative Sampling, which the paper cites as the single-SSM approach it extends §1 §3. Iteration-level scheduling and continuous batching come from Orca; the parallel layout comes from Megatron-LM §5.1. The attention kernel is adapted from FasterTransformer §5.2.
• Compared against: vLLM, HuggingFace Text Generation Inference and FasterTransformer for distributed serving §6.2; FlexGen for offloading §6.3; sequence-based speculative inference as an internal baseline §6.2.