Qwen3 moe mtp

[shihaobai]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @shihaobai, 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 comprehensive support for the Qwen3 Mixture-of-Experts (MoE) model within the Multi-Turn Prediction (MTP) framework. It involves defining a new Qwen3 MoE MTP model, adapting core components like FlashAttention and memory management to handle MTP-specific requirements, and updating the model loading mechanism to seamlessly integrate this new model type as a draft model. The changes ensure that the LightLLM system can efficiently utilize Qwen3 MoE models for improved multi-turn prediction performance.

Highlights

• Qwen3 MoE MTP Model Integration: Introduced a new model, 'Qwen3MOEMTPModel', extending 'Qwen3MOEModel' to support Multi-Turn Prediction (MTP) capabilities within the LightLLM framework.
• MTP-aware FlashAttention: Modified FlashAttention state initialization and attention calculations to dynamically adjust batch sizes and sequence lengths based on MTP steps, ensuring correct processing in multi-turn scenarios.
• Dynamic Memory Management for MTP: Enhanced memory manager initialization to account for additional layers required by draft models when operating in MTP mode, optimizing resource allocation.
• Draft Model Loading Extension: Updated the model inference backend to recognize and load 'Qwen3MOEMTPModel' as a valid MTP draft model, enabling its use in the prediction pipeline.

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 adds support for the Qwen3 MOE model in MTP (Multi-Token Prediction) mode. The changes include new model and weight classes for Qwen3MOEMTP and modifications to memory management, attention mechanisms, and model loading to accommodate this new model. While the changes are mostly correct, I've identified a few issues related to code clarity, maintainability, and a design flaw where model-specific logic is misplaced. My review includes suggestions to rename misleading variables, simplify redundant code, refactor hardcoded logic to a more appropriate location, and improve debugging messages.

[gemini-code-assist]

This logic for calculating added_mtp_layer_num is specific to deepseekv3 MTP modes, but it's located in the qwen2 model file. This creates a tight and incorrect coupling. When running with a Qwen MTP model, this logic will fail to calculate the correct number of extra layers for the memory manager, potentially leading to insufficient memory allocation and runtime errors. This logic should be generalized or moved to a more appropriate location, such as a base MTP model or handled during model initialization based on the specific draft model's configuration.

[gemini-code-assist]

The variable is_deepseekv3_mtp_draft_model is now misleading as it also checks for Qwen3MOEMTPModel. It should be renamed to something more generic, like is_mtp_draft_model, to accurately reflect its purpose.

Suggested change

	is_deepseekv3_mtp_draft_model = "Deepseek3MTPModel" in str(self.__class__) or "Qwen3MOEMTPModel" in str(
	self.__class__
	)
	if is_deepseekv3_mtp_draft_model:
	is_mtp_draft_model = "Deepseek3MTPModel" in str(self.__class__) or "Qwen3MOEMTPModel" in str(
	self.__class__
	)
	if is_mtp_draft_model:

[gemini-code-assist]

This if/else block can be simplified. The slicing logic [args_mtp_step :: (args_mtp_step + 1)] works correctly for args_mtp_step = 0 as well, where it becomes [0::1]. The .contiguous() call is necessary for slices with a step greater than 1 and is harmless for already contiguous tensors, so it can be applied unconditionally.

            self.b_att_seq_len = self.b_seq_len[args_mtp_step :: (args_mtp_step + 1)].contiguous()

[gemini-code-assist]

The assertion message in verify_load is not very descriptive. For better debuggability, it would be helpful to include the name of the weight that failed to load, rather than just its index.

Suggested change

	weights = [self.eh_proj_weight_, self.enorm_weight_, self.hnorm_weight_]
	for i in range(len(weights)):
	assert weights[i] is not None, "index:" + str(i) + " " + errors
	weight_names = ["eh_proj_weight_", "enorm_weight_", "hnorm_weight_"]
	weights = [self.eh_proj_weight_, self.enorm_weight_, self.hnorm_weight_]
	for i, name in enumerate(weight_names):
	assert weights[i] is not None, f"{name} {errors}"

[gemini-code-assist]

Reusing Deepseek3MTPPreLayerInfer for a Qwen model might be confusing due to its specific name. If this class contains generic MTP pre-layer inference logic, consider renaming it to something more abstract (e.g., MTPPreLayerInfer) to improve clarity and maintainability.