[Feature] enable index Cache for npu#8324
[Feature] enable index Cache for npu#8324ChefWu551 wants to merge 1 commit intovllm-project:mainfrom
Conversation
ChefWu551
requested review from
realliujiaxu,
weijinqian0,
whx-sjtu and
zzzzwwjj
as code owners
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
Summary of ChangesHello, I'm Gemini Code Assist1! 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 integrates IndexCache functionality into the vLLM-Ascend stack. By allowing attention layers to share and reuse top-k indices, the changes reduce redundant computations during model inference, aligning the NPU implementation with upstream vLLM improvements. Highlights
🧠 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. Using Gemini Code AssistThe 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
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 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. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a mechanism to share topk_indices across layers in AscendSFAMetadata using a skip_topk flag to optimize performance. A critical issue was identified where the shared indices are incorrectly reused despite being dependent on layer-specific hidden_states.
Suggested PR Title:
[Attention][Ops][Feature] Support shared top-k indices in SFASuggested PR Summary:
### What this PR does / why we need it?
This PR adds a `shared_topk_indices` field to `AscendSFAMetadata` and a `skip_topk` flag to the attention and MLA modules. This allows layers to reuse previously computed top-k indices to reduce redundant computations.
Feedback: A critical flaw was identified where `topk_indices` are computed from layer-specific `hidden_states`, making the reuse of these indices across layers mathematically incorrect.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
CI passed.| if self.skip_topk and attn_metadata.shared_topk_indices is not None: | ||
| topk_indices = attn_metadata.shared_topk_indices | ||
| else: | ||
| topk_indices = self.indexer_select_post_process( | ||
| x=hidden_states, | ||
| q_c=q_c, | ||
| kv_cache=kv_cache, | ||
| attn_metadata=attn_metadata, | ||
| cos=cos, | ||
| sin=sin, | ||
| actual_seq_lengths_query=actual_seq_lengths_query, | ||
| actual_seq_lengths_key=actual_seq_lengths_key, | ||
| ) | ||
| attn_metadata.shared_topk_indices = topk_indices |
There was a problem hiding this comment.
The current implementation of IndexCache appears to have a critical flaw. The topk_indices are cached and reused across layers, but their computation in indexer_select_post_process depends on hidden_states, which is unique to each layer.
Specifically, indexer_select_post_process uses x (which is hidden_states) to compute weights:
weights, _ = self.weights_proj(x)These weights are then used to determine topk_indices. Since hidden_states differ from one layer to the next, the topk_indices will also be different. Reusing them will lead to incorrect attention calculations.
For IndexCache to work correctly, the computation of topk_indices must be based on tensors that are shared across the layers intended to use the cache. This might require passing a shared tensor to indexer_select_post_process instead of the per-layer hidden_states.
There was a problem hiding this comment.
Thanks for the review. This behavior is intentional and matches upstream IndexCache semantics (vLLM PR #37735, issue #37684).
IndexCache is an approximate optimization: “full” layers compute top-k indices, while “shared” layers reuse cached indices to reduce redundant computation.
In our implementation, reuse is only enabled when skip_topk=True; otherwise indices are computed per layer as usual.
We’ll also attach accuracy/performance results to quantify the tradeoff.
ChefWu551
changed the title
|
This PR has been closed due to implementation issues. PR #8398 has fixed the corresponding functionality and provided relevant benchmark data, showing an improvement of 16%–18%. |
1 participant
Motivation
Implemented the corresponding NPU adaptation based on upstream IndexCache work.
This PR adds Ascend NPU adaptation for IndexCache in vLLM-Ascend, based on:
Modifications
This PR includes NPU-oriented integration and adaptation for IndexCache in vLLM-Ascend
Accuracy Tests
Benchmarking and Profiling
Checklist
Format your code according to the Format code with pre-commit.
Update documentation according to Write documentations. (if required)
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. (pending final data upload)
Follow the SGLang code style guidance.
vLLM version:
vLLM main: https://github.com/vllm-project/vllm/commit/v0.19.0