Don't fuse wk+weight_proj for nextn

[trevor-m]

Motivation

Potential fix for #12841

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

[gemini-code-assist]

Summary of Changes

Hello @trevor-m, 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 introduces a targeted modification to the deepseek_v2 model's weight loading mechanism. Specifically, it disables the fusion of wk and weights_proj when the is_nextn flag is active. This adjustment is intended to address a reported bug by ensuring proper weight management under specific operational conditions, thereby improving model stability or correctness.

Highlights

• Conditional Weight Fusion: Prevents the fusion of wk and weights_proj in the deepseek_v2 model's load_weights function when is_nextn is enabled.
• Bug Fix: This change is a potential fix for an identified issue related to model weight handling (issue [Bug] Deepseek-V3.2-EXP FP4 checkpoint fails to load when MTP is enabled #12841).

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 pull request correctly disables the wk and weight_proj fusion for nextn models by adding and not is_nextn to the condition. This appears to be a good fix for the underlying issue. I have one minor suggestion to improve the readability of the code.

[gemini-code-assist]

The logic is correct. For better readability, the boolean expression can be enclosed in a single set of parentheses. Additionally, the comment can be made more concise to focus on the specific reason for the change.

# For nextn, fp4 is disabled, so we cannot fuse wk and weights_proj.
fuse_wk_and_weights_proj = (is_nsa_indexer_wk_and_weights_proj_fused(
    self.config, self.quant_config
) and not is_nextn)