[Refact]Refact MLA/SFA weight prefetch to consist with moe weight prefetch

[leo-pony]

What this PR does / why we need it?

1. [Refact] Refact MLA/SFA weight prefetch to consist with moe weight prefetch
2. Remove duplicated o_proj weight prefetch in forward for MLA/SFA

Does this PR introduce any user-facing change?

NA

How was this patch tested?

1. Performance result:
   Perf test data:
   *) MLA:

	1st test	2nd test	Output Token Throughput(Avg)	Performance improvement percentage
o_proj duplicate prefetch	11.9669 token/s	12.0287 token/s	11.9978
o_proj no duplicate prefetch	12.5594 token/s	12.6216 token/s	12.5905	4.94%

single layer performace improve: 5%~8%

*) SFA:

	1st test	2nd test	Output Token Throughput(Avg)	Performance improvement percentage
o_proj duplicate prefetch	13.0523 token/s	13.1084 token/s	13.08035
o_proj no duplicate prefetch	13.9844 token/s	14.1678 token/s	14.0761	7.6%

2. Accuracy test result:

3. Performance profiling record
   Profiling duplicated o_proj weight prefetch been removed:
   *) before:

*) After, duplicated o_proj weight prefetch has been remove:

• vLLM version: v0.15.0
• vLLM main: vllm-project/vllm@d7e17aa

[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]

Summary of Changes

Hello @leo-pony, 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 significantly refactors the weight prefetching architecture for MLA and SFA attention modules, bringing their implementation in line with the MoE weight prefetching strategy. The changes aim to streamline the prefetching logic by centralizing control and introducing a more robust, unified method for handling weight prefetch operations. This refactoring also includes specific optimizations, such as conditionally skipping prefetching for certain quantized weights, to improve overall efficiency and maintainability.

Highlights

• Unified Weight Prefetching: Refactored MLA (Multi-Layer Attention) and SFA (Self-Attention) weight prefetching mechanisms to align with the existing MoE (Mixture of Experts) weight prefetch logic, centralizing control within the WeightPrefetchMethod class.
• Centralized Prefetch Control: Removed individual enable_prefetch flags from MLAImpl and SFAImpl classes, consolidating the prefetch enablement into a single other_prefetch_enable attribute within WeightPrefetchMethod.
• New Prefetch Method: Introduced a new method, maybe_prefetch_weight_in_current_stream, within WeightPrefetchMethod to encapsulate the general weight prefetching logic, replacing direct calls to the older maybe_npu_prefetch function.
• Quantization-Aware Prefetching: Added a condition to skip output projection (o_proj) weight prefetching for models using AscendW8A8LinearMethod quantization in both MLA and SFA modules, optimizing performance for these specific configurations.
• Improved MoE Prefetch Robustness: Enhanced the maybe_prefetch_moe_weight_preprocess method with null checks for forward_context and model_instance to prevent potential errors during MoE weight prefetching.

🧠 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/attention/mla_v1.py
  • Updated imports to use get_weight_prefetch_method from vllm_ascend.utils and removed the import of maybe_npu_prefetch.
  • Removed the self.enable_prefetch attribute from the MLAImpl class constructor.
  • Replaced direct calls to maybe_npu_prefetch with weight_prefetch_method.maybe_prefetch_weight_in_current_stream for fused_qkv_a_proj and o_proj weights.
  • Added a conditional check to prevent o_proj weight prefetching when using AscendW8A8LinearMethod quantization.
• vllm_ascend/attention/sfa_v1.py
  • Updated imports to use get_weight_prefetch_method from vllm_ascend.utils and removed the import of maybe_npu_prefetch.
  • Removed the self.enable_prefetch attribute from the SFAImpl class constructor.
  • Replaced direct calls to maybe_npu_prefetch with weight_prefetch_method.maybe_prefetch_weight_in_current_stream for fused_qkv_a_proj and o_proj weights.
  • Added a conditional check to prevent o_proj weight prefetching when using AscendW8A8LinearMethod quantization.
• vllm_ascend/ops/weight_prefetch.py
  • Added self.other_prefetch_enable attribute to the WeightPrefetchMethod constructor to centralize the enable/disable flag for non-MoE prefetching.
  • Introduced null checks for forward_context and forward_context.model_instance in maybe_prefetch_moe_weight_preprocess for increased robustness.
  • Implemented a new method maybe_prefetch_weight_in_current_stream within WeightPrefetchMethod to handle general weight prefetching logic, utilizing the other_prefetch_enable flag.

Activity

• The pull request was initiated by leo-pony.
• All Continuous Integration (CI) tests have passed, indicating the changes are functional and do not introduce regressions.

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. ↩

[gemini-code-assist]

Code Review

This is a good refactoring that unifies the weight prefetch logic for MLA/SFA with the MoE implementation, improving code consistency and maintainability. I have one comment regarding a redundant import. Additionally, I've provided suggestions for the PR title and summary to align with the project's style guide.

Suggested PR Title:

[Attention][Refact] Unify MLA/SFA weight prefetch with MoE prefetch

Suggested PR Summary:

### What this PR does / why we need it?
This pull request refactors the weight prefetching mechanism for Multi-Head Latent Attention (MLA) and Sparse Flash Attention (SFA) to align with the implementation used for Mixture-of-Experts (MoE) models.

Key changes include:
- Centralizing the prefetching logic by introducing a new method `maybe_prefetch_weight_in_current_stream` in the `WeightPrefetchMethod` class.
- Replacing direct calls to `maybe_npu_prefetch` in `mla_v1.py` and `sfa_v1.py` with the new unified method.
- Adding an optimization to skip prefetching for `o_proj` weights when W8A8 quantization is used.
- Adding a defensive check in `maybe_prefetch_moe_weight_preprocess` to prevent potential `None` access errors.

These changes improve code consistency and maintainability by reducing code duplication and creating a single, unified path for weight prefetching.

### Does this PR introduce _any_ user-facing change?
NA

### How was this patch tested?
- vLLM version: v0.15.0
- vLLM main: https://github.com/vllm-project/vllm/commit/d7e17aaacd5ed1b4b4be6bcfef3a1b7cbc84fc9a