[MTP][BugFix] Preserve MoE weight loaders during online updates

[Mengyuyang]

What this PR does / why we need it?

This PR prevents existing weight_loader metadata from being dropped when unquantized Ascend MoE weights change layout during online RL weight updates with MTP.

The target model and the MTP drafter receive actor weights through model.load_weights(). After loading, Ascend converts w13_weight and w2_weight from the loadable layout to the execution layout. The previous implementation registered raw torch.nn.Parameter objects during this conversion, which dropped the existing expert-aware weight_loader.

In addition, NPUWorker.wake_up() restored the loadable layout only for the target model, while the MTP drafter is a separate model that receives the same online updates.

This change:

• uses vLLM's public replace_parameter() helper when converting unquantized MoE weights in both layout directions, preserving an existing weight_loader;
• prepares both the target model and the MTP drafter during wake_up() when speculative_config.method == "mtp";
• keeps non-MTP drafter behavior unchanged;
• adds unit coverage for loader preservation, target/drafter wake-up handling, and non-MTP isolation.

The change does not add model code, environment variables, global state, or a framework-side patch.

Does this PR introduce any user-facing change?

Yes. In MTP mode, NPUWorker.wake_up() now applies the existing unquantized MoE layout preparation to both the target and the MTP drafter.

There is no API or configuration change, and non-MTP speculative methods keep their existing behavior.

How was this patch tested?

Unit coverage is included for:

• preserving w13_weight and w2_weight layout, contiguity, and weight_loader after process_weights_after_loading();
• preparing both target and MTP drafter weights in NPUWorker.wake_up();
• confirming that non-MTP drafters are not included in the new behavior.

Manual integration validation was performed with:

• Ascend NPU;
• Qwen3.5-35B-A3B;
• MTP speculative decoding enabled;
• repeated online actor-to-rollout weight updates;
• graph rollout modes.

Both modes completed multiple training iterations without weight_loader errors. Draft acceptance and reward remained normal.

Additional checks:

git diff --check
python -m py_compile \
  tests/ut/ops/test_fused_moe.py \
  tests/ut/worker/a2/test_worker_v1.py \
  vllm_ascend/ops/fused_moe/routed_experts.py \
  vllm_ascend/worker/worker.py
- vLLM version: v0.26.0
- vLLM main: https://github.com/vllm-project/vllm/commit/d02df748bf9efd99022f1a062597dc3cb3808485

[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 addresses a bug where MoE weight metadata was inadvertently dropped during online weight updates on Ascend NPUs. By switching to a parameter replacement utility that respects existing metadata, the system now correctly maintains expert-aware loaders. Additionally, the wake-up logic has been extended to ensure that MTP drafter models receive the same necessary layout preparations as the primary target model, ensuring consistency during online RL weight updates.

Highlights

• Weight Loader Preservation: Replaced direct torch.nn.Parameter instantiation with vLLM's replace_parameter helper to ensure existing weight_loader metadata is preserved during MoE weight layout conversions.
• MTP Drafter Support: Updated NPUWorker.wake_up() to apply weight layout preparation to both the target model and the MTP drafter model when MTP speculative decoding is enabled.
• Testing Coverage: Added comprehensive unit tests to verify loader preservation, correct target/drafter wake-up handling, and isolation of non-MTP speculative methods.

New Features

🧠 You can now enable Memory (public preview) 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 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 the 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 counterproductive. 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

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.

---

Tip

💡 Consider Linking a Related Issue or RFC

Your PR title contains the [BugFix] tag, indicating a bug fix or new feature.

Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:

• Fixes #<issue_number>
• Closes #<issue_number>
• Resolves #<issue_number>
• Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

[gemini-code-assist]

Code Review

This pull request ensures that weight_loader metadata is preserved during online weight updates for unquantized Ascend MoE weights by utilizing vLLM's replace_parameter helper. It also updates NPUWorker.wake_up() to prepare both the target model and the MTP drafter model when speculative decoding is enabled with the 'mtp' method. The review feedback correctly identifies a style guide violation regarding the PR title format and provides a compliant title and summary suggestion.

[gemini-code-assist]

The Pull Request Title and Summary do not fully adhere to the Repository Style Guide. Specifically, the title format should be [Branch][Module][Action] Pull Request Title. The current title uses [BugFix][MTP] which swaps the Action and Module.\n\nHere are the suggested PR Title and Summary according to the style guide:\n\nSuggested PR Title:\n\nmarkdown\n[MTP][BugFix] Preserve MoE weight loaders during online updates\n\n\nSuggested PR Summary:\n\nmarkdown\n### What this PR does / why we need it?\n\nThis PR prevents existing `weight_loader` metadata from being dropped when unquantized Ascend MoE weights change layout during online RL weight updates with MTP.\n\nThe target model and the MTP drafter receive actor weights through `model.load_weights()`. After loading, Ascend converts `w13_weight` and `w2_weight` from the loadable layout to the execution layout. The previous implementation registered raw `torch.nn.Parameter` objects during this conversion, which dropped the existing expert-aware `weight_loader`.\n\nIn addition, `NPUWorker.wake_up()` restored the loadable layout only for the target model, while the MTP drafter is a separate model that receives the same online updates.\n\nThis change:\n- Uses vLLM's public `replace_parameter()` helper when converting unquantized MoE weights in both layout directions, preserving an existing `weight_loader`.\n- Prepares both the target model and the MTP drafter during `wake_up()` when `speculative_config.method == "mtp"`.\n- Keeps non-MTP drafter behavior unchanged.\n- Adds unit coverage for loader preservation, target/drafter wake-up handling, and non-MTP isolation.\n\n### Does this PR introduce _any_ user-facing change?\n\nYes. In MTP mode, `NPUWorker.wake_up()` now applies the existing unquantized MoE layout preparation to both the target and the MTP drafter.\n\nThere is no API or configuration change, and non-MTP speculative methods keep their existing behavior.\n\n### How was this patch tested?\n\nUnit coverage is included for:\n- Preserving `w13_weight` and `w2_weight` layout, contiguity, and `weight_loader` after `process_weights_after_loading()`.\n- Preparing both target and MTP drafter weights in `NPUWorker.wake_up()`.\n- Confirming that non-MTP drafters are not included in the new behavior.\n\nManual integration validation was performed with:\n- Ascend NPU;\n- Qwen3.5-35B-A3B;\n- MTP speculative decoding enabled;\n- repeated online actor-to-rollout weight updates;\n- graph rollout modes.\n

References

1. The Pull Request Title and Summary must follow the specified format in the Repository Style Guide. The current title '[BugFix][MTP]' has the Action and Module swapped (should be '[MTP][BugFix]'). ^(link)