Speculative decoding on sglang

• Slide

What is Tree-Attention?

• Problem statements - Existing Methods (basic speculative decoding)

  

  • incremental decoding (auto-regressive decoding)
    • Generates one token at a time, which is highly inefficient for data reuse. Consequently, LLM decoding operations are known to be heavily memory-bound.
  • Sequence-based Speculative Decoding

    • To address the memory-bound issue, a small draft model generates multiple (gamma) tokens in advance, and the target LLM verifies these draft tokens in parallel.

      

    • The major issue with the existing sequence-based speculative decoding is that if even a single token in the sequence is rejected, all subsequent tokens are discarded.

    • Example)

      

    • To address this issue, a new technique called "Tree attention-based Speculative Decoding" was introduced.

• Tree-based Speculative Decoding

  • https://arxiv.org/abs/2305.09781 (Tree-based Speculative Decoding 처음 제안한 논문)

  

  • Unlike the traditional sequence-based method, tree-attention generates draft tokens in a tree structure and passes them through the target model in parallel for verification.

  • While the target model still performs parallel verification, the key advantage is that multiple candidate trees increase the acceptance rate.

  • However, since the target model verifies tokens in parallel, proper handling of the tree_mask is essential.

  • tree_mask , tree_structure Example

    

Tree-based Speculative Decoding (EAGLE) Implementation Summary

1. Prefill: This step is the same as standard auto-regressive decoding in LLMs.
2. Generate Draft Tree Candidates: Given an input prompt, the draft model (SLM) generates draft tokens for gamma time steps.
3. Tree Decoding: The generated draft candidates are flattened into a 1-dimensional input_ids format. Along with this, both attention_mask and tree_attention_mask are used to ensure that attention operations are skipped for unrelated draft tokens.

As a result, the target model determines which tokens are verified (accepted) and updates the output accordingly.

( Important Implementation Detail*: It is necessary to selectively update the KV Cache based on the indices of accepted tokens.)

Issue: Speculative Decoding on vLLM

• Currently, vLLM only implements a top-1 proposer (i.e., sequence-based) engine.

  • Sequence-based speculative decoding is a relatively naive approach.
    • If rejected tokens occur, theoretically, it can function without issues by simply modifying the cache_position or position_ids for the corresponding indices in the next decoding step.
  • However, tree-based attention necessarily requires a KV cache rewrite process, making this step crucial.

• Reference: https://github.com/vllm-project/vllm/blob/256a2d29dc2358d7c0a5d38c0faf152095335929/vllm/spec_decode/spec_decode_worker.py#L132

• Tree-attention is left as future work (No significant progress for over six months).