[V1][Spec Decode] Support multi-layer eagle draft model

[zixi-qi]

Eagle draft model with > 1 layers would fail after #17394

This PR removes this assertion and let all draft layers share the same AttentionMetadata (if I understand this correctly: https://github.com/heheda12345/vllm/blob/77d86dd096878fbaa834ab5411720b0e1b160da9/vllm/v1/worker/gpu_model_runner.py#L1256 )

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[zixi-qi]

cc @heheda12345 please let me know if this makes sense to you, this PR works for us but from #17394 I assume the ideal way to do this is to create a kv cache group to be used by draft model?

[heheda12345]

The eagle layers do not be put into a special kv cache group. I think we can support multi-layer eagle by building a dict[layer_name, AttentionMetadata] here.

vllm/vllm/v1/spec_decode/eagle.py

Line 137 in f6518b2

with set_forward_context(attn_metadata,

[zixi-qi]

The eagle layers do not be put into a special kv cache group. I think we can support multi-layer eagle by building a dict[layer_name, AttentionMetadata] here.

vllm/vllm/v1/spec_decode/eagle.py

Line 137 in f6518b2

with set_forward_context(attn_metadata,

The dict[layer_name, AttentionMetadata] for each eagle layer can already be supported out of the box, I think we just need to find a way to propagate the per layer AttentionMetadata to eagle layers instead of
https://github.com/heheda12345/vllm/blob/77d86dd096878fbaa834ab5411720b0e1b160da9/vllm/v1/worker/gpu_model_runner.py#L1256

Would you have any suggestions based on your new design? Thanks in advance

[heheda12345]

For now, I think we can add an assertion to verify that all eagle layers are in the same kv cache group in GPUModelRunner.initialize_kv_cache so that they use the same attention metadata. This assumption is true if eagle part only contains full attention layers and does not contain sparse attention like sliding window attention and I think it is enough now.
In the long term, we may need to build different eagle attention metadata for different layers.

And can you help to change the code here to dict[layer_name, FlashAttentionMetadata] by building one FlashAttentionMetadata and assign it to all layers?

vllm/vllm/v1/spec_decode/eagle.py

Line 113 in 0b217da

attn_metadata = FlashAttentionMetadata(

[zixi-qi]

For now, I think we can add an assertion to verify that all eagle layers are in the same kv cache group in GPUModelRunner.initialize_kv_cache so that they use the same attention metadata. This assumption is true if eagle part only contains full attention layers and does not contain sparse attention like sliding window attention and I think it is enough now. In the long term, we may need to build different eagle attention metadata for different layers.

And can you help to change the code here to dict[layer_name, FlashAttentionMetadata] by building one FlashAttentionMetadata and assign it to all layers?

vllm/vllm/v1/spec_decode/eagle.py

Line 113 in 0b217da

attn_metadata = FlashAttentionMetadata(

Updated, please take a look when you get a chance

[heheda12345]

Thanks for your contribution. It aligns to the design in general. Left some small comments. And can we add a test for multi-layer eagle?

[zixi-qi]

Thanks for your contribution. It aligns to the design in general. Left some small comments. And can we add a test for multi-layer eagle?

Thanks for the review and detailed feedback! For testing AFAIK there isn't a open source multi-layer eagle model yet to be used in tests, we will add the test once we release ours

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zixi-qi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[heheda12345]

LGTM. Only some nits.

[heheda12345]

Suggested change

	def validate_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
	def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
	"""
	Validate that all eagle layers belong to the same KVCacheGroup. Need this assumption to ensure all eagle layers can use the same AttentionMetadata. May extend to multiple AttentionMetadata and remove this assumption in the future.
	"""

[zixi-qi]

@heheda12345 Thanks for adding the comment, updated them in the PR. With regards to the helper function naming, I prefer to keep it more generic as 1. the comments have enough details 2. we may want to do other types of kv cache validation for eagle model in the future

[heheda12345]

What further kv cache validation do we need for eagle model?

[zixi-qi]

Depending on future development on the eagle model, we may need to do further validation. I don't think we would know right now.

To take a step back, I think helper function naming is relatively a trivial problem. Given all other changes are aligned, I would suggest that we move forward to merge this PR so that the issue with multi-layer eagle model affecting our use case can be properly fixed. WDYT? cc @luccafong @WoosukKwon

[heheda12345]

I think we should always make the function name as clear as possible. And future development should remove such constraint to improve compatibility, instead of adding more constraint.

[zixi-qi]

Thanks again for sharing your opinion. Agree that it would be ideal that such constraint can be removed in the future. Though I think it would really depend of the future development direction of eagle model architecture, which is hard to tell right now.

For function names, with the amount of clarity that we already have thanks to your added comments, I think it comes down to a personal preference nit rather a blocking issue. Let's get a review from community to fully resolve this :-)

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zixi-qi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[WoosukKwon]

@zixi-qi Could you please rebase the PR and resolve the merge conflicts?

@heheda12345 Could you please take the final pass?

[heheda12345]

@WoosukKwon My only concern is the function name discussed here https://github.com/vllm-project/vllm/pull/18030/files#r2094177259.
But @zixi-qi wants other people instead of me to review.

I can do a final pass after this problem is resolved & the rebase.

[zixi-qi]

@WoosukKwon My only concern is the function name discussed here https://github.com/vllm-project/vllm/pull/18030/files#r2094177259. But @zixi-qi wants other people instead of me to review.

I can do a final pass after this problem is resolved & the rebase.

@heheda12345 Changed the helper function name from "validate_kv_cache_group" to "validate_same_kv_cache_group" as you requested. Thanks for the discussion.

cc @WoosukKwon

[heheda12345]

LGTM! Thanks for your contribution.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zixi-qi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork