[Feat] Support separate attention backend for target and draft model.

[SidaoY]

What this PR does / why we need it?

This PR enables separate attention backend configuration for target and draft models in speculative decoding, decoupling the previously bound attention backend settings between the two models.

It solves the compatibility issue where some draft models do not support the attention backend used by the target model, and allows users to select the optimal attention backend for each model individually to maximize inference performance. The change is fully backward compatible.

Does this PR introduce any user-facing change?

• Full backward compatibility: existing global attention backend configurations work without any modifications.
• New feature: support specifying a dedicated attention backend for the draft model for fine-grained configuration in speculative decoding.

How was this patch tested?

• Tested on vLLM v0.17.0 (upstream commit 4034c3d)
• Backward compatibility test: verified no regression with original global backend configuration
• Functional test: validated speculative decoding works correctly with different attention backend combinations for target and draft models
• Model compatibility test: covered mainstream LLM + draft model combinations
• End-to-end test: confirmed correct inference output and performance improvement with optimal backend selection

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the KV cache allocation and reshaping mechanisms within the model runner. The primary goal is to enhance flexibility and correctness when handling different model architectures, particularly those with varying attention layer configurations (e.g., MLAAttention). By introducing helper methods to retrieve layer-specific KV cache specifications and dimensions, the system can now accurately allocate and reshape KV cache tensors tailored to each layer's unique requirements, improving compatibility and efficiency across diverse target and draft models.

Highlights

• KV Cache Specification Retrieval: Introduced a new helper method, _get_layer_kv_cache_specs, to efficiently retrieve and map layer-specific KV cache configurations, supporting both uniform and varied specifications across different attention groups.
• Dynamic Attention KV Cache Dimension Calculation: Added _get_attention_kv_cache_dims to dynamically determine the K and V dimensions for attention layers, specifically handling MLAAttentionSpec by extracting rank and head dimensions from the attention layer itself, and using head sizes for other attention types.
• Refactored KV Cache Allocation and Reshaping: Updated the KV cache allocation (_allocate_kv_cache_tensors) and reshaping (_reshape_kv_cache_tensors) logic to leverage the new layer-specific KV cache specifications and dynamic dimension calculations, ensuring accurate sizing and formatting for diverse model architectures, including those with MLAAttention and sparse attention.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• vllm_ascend/worker/model_runner_v1.py
  • Added _get_layer_kv_cache_specs method to map layer names to their specific KV cache specifications.
  • Added _get_attention_kv_cache_dims method to calculate K and V dimensions based on attention layer type.
  • Updated _allocate_kv_cache_tensors to utilize the new _get_layer_kv_cache_specs for retrieving layer-specific KV cache configurations and to dynamically calculate K/V tensor split factors based on attention layer dimensions.
  • Modified _reshape_kv_cache_tensors to consistently use layer-specific KV cache specifications for block size, head dimensions, page size, and data types, ensuring accurate tensor reshaping for various attention and Mamba layer types.

Activity

• No specific activity has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Code Review

This pull request refactors the KV cache allocation logic to correctly handle different attention specifications, particularly for models with UniformTypeKVCacheSpecs and MLAAttentionSpec. The changes extract helper methods to get layer-specific cache specs and attention dimensions, making the allocation and reshaping of KV cache tensors more generic and robust. This fixes bugs where incorrect cache specs were used for layers within a group.

However, I've identified a critical issue where a division by zero can occur during the calculation of KV cache split factors if k_dim or v_dim is zero, which can happen with certain MLAAttentionSpec configurations. I've provided a suggestion to prevent this crash.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.