[Feature]: Tree-Attention Support for Speculative Decoding

[yesredpig]

🚀 The feature, motivation and pitch

Summary

We propose adding tree-attention-based speculative decoding support to vLLM to improve inference throughput, token acceptance rates, and memory efficiency during generation. This approach is inspired by SpecInfer ([arXiv:2305.09781](https://arxiv.org/abs/2310.08560)) and EAGLE2 (arXiv:2406.16858), both of which demonstrate state-of-the-art performance using token tree structures and topology-aware attention for parallel decoding.

---

Motivation

Current speculative decoding strategies in vLLM rely on batch expansion or multi-head proposals. These approaches face key limitations:

• ❌ Low token acceptance rates, especially in long sequences or large models
• ❌ Redundant computation and memory traffic due to duplicate KV cache updates
• ❌ Inefficient parallelism across speculative paths

Both SpecInfer and EAGLE2 propose using tree-based token structures combined with topology-aware masking, enabling efficient evaluation of multiple speculative paths in a single kernel invocation. This approach improves both accuracy and decoding speed.

---

Reference: SpecInfer and EAGLE2

🔍 SpecInfer (ASPOS’24)

• Introduces a tree-structured speculative decoding algorithm
• Uses topo-aware attention masking and single-kernel tree scoring
• Achieves higher token throughput and better acceptance across large LLMs

Throughput and latency comparison between SpecInfer and standard speculative decoding:

• Up to 3× throughput improvement
• Significant latency reduction per generated token
• Works efficiently on long sequences and large target models

🔍 EAGLE2 (Empirical Methods in Natural Language Processing’24)

EAGLE2 is another recent method using tree attention to accelerate speculative decoding. It builds a hierarchical token tree and applies a mask-aware attention mechanism to support efficient tree traversal and scoring. The algorithm enables parallel generation of token drafts in a tree structure and shows competitive performance for decoding large models.

EAGLE-2 is:

• 4x faster than vanilla decoding (13B).
• 1.4x faster than EAGLE-1 (13B).

---

Expected Benefits

✅ Significantly increased decoding throughput (tokens/sec)
✅ Improved token acceptance rates from the target model
✅ Lower HBM usage via shared prefix reuse
✅ Strong scalability with large sequence lengths and large model sizes

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[yesredpig]

If our team PL explains Implementation, it will be deleted.

---

Proposal

We propose foundational support for tree-attention speculative decoding in vLLM. This includes:

1. Token Tree Proposer

• Builds a token tree with configurable depth and width.
• Generates topology-aware attention masks to represent valid branches for scoring.

2. Tree-Aware Scorer

• Modify the FlashAttention var-length kernel to support topology-aware attention.
• Enables parallel scoring of the full tree in a single kernel call.
• Make TreeAttetionScorer with FlashAttention kernel, (currently SD worked MQA, BatchExpension Scorer )

3. Token Tree Management

• Tracks token tree metadata: parent-child links, path indices, and accepted branches.
• Supports batched inference and dynamic tree expansion. with dynamic K

4. Runtime API

• Update speculative decoding API to support tree-based proposals.
• Maintain backward compatibility with flat proposals (single path).

5. KV Cache

• prefix reuse (reference sharing), deta KV caching

---

Out-of-Scope for Now

• ❌ Multi-SSM support
• ❌ Tensor Parallelism for SSM Proposers

These limitations remain and are not addressed in this initial request.

---

Next Steps

This feature request is the first in a series to enable full tree-attention speculative decoding in vLLM. Once the foundational elements are in place, further enhancements such as dynamic branching and SSM-tree co-design can be layered on.
We welcome design feedback and are happy to align with internal refactoring plans or roadmap direction.

[wwl2755]

Hi @yesredpig ,
Thanks for bringing up this issue. Is this feature already done within your team or ongoing development?

I'm currently working on the tree-attention under hood of EAGLE. (Ref: Issue, PR and Doc)

[wwl2755]

If our team PL explains Implementation, it will be deleted.

Proposal

We propose foundational support for tree-attention speculative decoding in vLLM. This includes:

1. Token Tree Proposer

• Builds a token tree with configurable depth and width.
• Generates topology-aware attention masks to represent valid branches for scoring.

2. Tree-Aware Scorer

• Modify the FlashAttention var-length kernel to support topology-aware attention.
• Enables parallel scoring of the full tree in a single kernel call.

3. Token Tree Management

• Tracks token tree metadata: parent-child links, path indices, and accepted branches.
• Supports batched inference and dynamic tree expansion.

4. Extended Proposal and Score APIs

• Update speculative decoding API to support tree-based proposals.
• Maintain backward compatibility with flat proposals (single path).

Out-of-Scope for Now

• ❌ Multi-SSM support
• ❌ Tensor Parallelism for SSM Proposers

These limitations remain and are not addressed in this initial request.

Next Steps

This feature request is the first in a series to enable full tree-attention speculative decoding in vLLM. Once the foundational elements are in place, further enhancements such as dynamic branching and SSM-tree co-design can be layered on. We welcome design feedback and are happy to align with internal refactoring plans or roadmap direction.

I believe the proposals are quite similar. If you feel comfortable and interested, we could break it to multiple steps and collaborate on this.

[yesredpig]

If our team PL explains Implementation, it will be deleted.

Proposal

We propose foundational support for tree-attention speculative decoding in vLLM. This includes:
1. Token Tree Proposer

• Builds a token tree with configurable depth and width.
• Generates topology-aware attention masks to represent valid branches for scoring.

2. Tree-Aware Scorer

• Modify the FlashAttention var-length kernel to support topology-aware attention.
• Enables parallel scoring of the full tree in a single kernel call.

3. Token Tree Management

• Tracks token tree metadata: parent-child links, path indices, and accepted branches.
• Supports batched inference and dynamic tree expansion.

4. Extended Proposal and Score APIs

• Update speculative decoding API to support tree-based proposals.
• Maintain backward compatibility with flat proposals (single path).

Out-of-Scope for Now

• ❌ Multi-SSM support
• ❌ Tensor Parallelism for SSM Proposers

These limitations remain and are not addressed in this initial request.

Next Steps

This feature request is the first in a series to enable full tree-attention speculative decoding in vLLM. Once the foundational elements are in place, further enhancements such as dynamic branching and SSM-tree co-design can be layered on. We welcome design feedback and are happy to align with internal refactoring plans or roadmap direction.

I believe the proposals are quite similar. If you feel comfortable and interested, we could break it to multiple steps and collaborate on this.

@wwl2755 Thanks for reply.
We are currently deciding on the design architecture and developing it. We are developing the topK selector in the TreeProposer and the flash-attention backend. We will share our design architecture and development details soon.

Basically, the part of developing TreeProposer seems to be sharable.
However, However, EAGLE2 explicitly creates the tree structure (for EAGLE2 model), so it is different from the dynamicTree implementation we are aiming for. We are considering TreeScorer, TreeRejection sampler, and KV handling for tree structure processing.

I will contact you personally.

[longer-is-better]

@wwl2755 Thanks for reply. We are currently deciding on the design architecture and developing it. We are developing the topK selector in the TreeProposer and the flash-attention backend. We will share our design architecture and development details soon.

Basically, the part of developing TreeProposer seems to be sharable. However, However, EAGLE2 explicitly creates the tree structure (for EAGLE2 model), so it is different from the dynamicTree implementation we are aiming for. We are considering TreeScorer, TreeRejection sampler, and KV handling for tree structure processing.

I will contact you personally.

Could you please give more details about Tree Scorer, Tree Rejection sampler, and KV handling for tree structure processing. Or recommend some related papers? @yesredpig

[yesredpig]

@wwl2755 Thanks for reply. We are currently deciding on the design architecture and developing it. We are developing the topK selector in the TreeProposer and the flash-attention backend. We will share our design architecture and development details soon.
Basically, the part of developing TreeProposer seems to be sharable. However, However, EAGLE2 explicitly creates the tree structure (for EAGLE2 model), so it is different from the dynamicTree implementation we are aiming for. We are considering TreeScorer, TreeRejection sampler, and KV handling for tree structure processing.
I will contact you personally.

Could you please give more details about Tree Scorer, Tree Rejection sampler, and KV handling for tree structure processing. Or recommend some related papers? @yesredpig

We are currently implementing TreeProposer, TreeScorer, and TreeRejection sampler on vLLM. The paper and code will be released after securing IP rights, so we cannot disclose all the details at this time.

For related references, please refer to

• SpecInfer (https://arxiv.org/abs/2305.09781) and
• FlexFlow Serve, which is based on SpecInfer (https://github.com/flexflow/flexflow-serve).

If further discussion is needed, please feel free to reach out.