Implement MTP for MiniCPM Model by yhyang201 · Pull Request #6899 · sgl-project/sglang

yhyang201 · 2025-06-05T15:04:47Z

Motivation

Implement MTP for MiniCPM Model.

Unit tests and documentation will be added after the relevant content has been released.

Many thanks to @mickqian and the @OpenBMB Team for their support.

python -m sglang.test.send_one

w/ mtp
acc_length=1.71
speed=136.74 token/s

---

w/o mtp
acc_length=1.00
speed=109.91 token/s

mmlu

w/ mtp
Total latency: 67.199
Average accuracy: 0.715

w/o mtp
Total latency: 61.985
Average accuracy: 0.715

Modifications

Checklist

Format your code according to the Code Formatting with Pre-Commit.
Add unit tests as outlined in the Running Unit Tests.
Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

gemini-code-assist

Hello @yhyang201, 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!

Summary of Changes

Hello! Gemini here, providing a summary of this pull request to help everyone quickly understand the changes.

This PR introduces support for the Multi-Token Prediction (MTP) variant of the MiniCPM model within the SGLang framework. This is a significant addition, likely enabling speculative decoding or similar techniques by allowing a smaller, faster MiniCPM model to propose multiple tokens ahead of the main model.

The core changes involve adding the specific MTP model implementation and updating the configuration to recognize MiniCPM as a potential draft model for MTP.

Highlights

Configuration Update: Modifies the model configuration logic to correctly identify the MiniCPM model as an MTP draft model when specified.
Base MiniCPM Model Adjustments: Makes minor adjustments to the base MiniCPM model class, including adding a method to retrieve embedding and language model head weights and removing an explicit setting for the rotary embedding cache precision.

Changelog

Click here to see the changelog

python/sglang/srt/configs/model_config.py
- Added logic to recognize 'MiniCPMForCausalLM' as 'MiniCPMMTP' when is_draft_model is true (lines 104-105).
python/sglang/srt/models/minicpm.py
- Removed explicit setting of rotary_emb.cos_sin_cache to fp32 (lines 141-142).
- Added get_embed_and_head method to retrieve embedding and lm_head weights (lines 336-337).
python/sglang/srt/models/minicpm_mtp.py
- Added a new file implementing the MiniCPMMTP model.
- Defined MiniCPMMultiTokenPredictorLayer for MTP-specific forward pass logic.
- Defined the main MiniCPMMTP class wrapping the MTP layer, parallel LM head, and logits processor.
- Implemented load_weights method with specific handling for stacked parameters and skipping embedding weights.
- Added get_embed_and_head and set_embed_and_head methods for weight sharing.
- Set EntryClass to MiniCPMMTP.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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.

Tokens fly so fast,
Predicting many at once,
Speedy MiniCPM.

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

The pull request implements MTP (Multi-Token Prediction) for the MiniCPM model. This introduces a new model architecture (MiniCPMMTP) and modifies existing configurations to support it. The changes seem well-structured, but let's review for potential improvements and ensure correctness.

Summary of Findings

Missing Comment: A comment is missing explaining why a line was commented out in python/sglang/srt/models/minicpm.py.
Missing Docstring: A docstring is missing for the get_embed_and_head function in python/sglang/srt/models/minicpm.py.
Unnecessary Print Statements: Print statements are present in the load_weights function of python/sglang/srt/models/minicpm_mtp.py that should be removed.

Merge Readiness

The pull request introduces MTP support for MiniCPM, which is a valuable addition. However, there are a few issues that should be addressed before merging, particularly the missing comment and docstring, and the commented out line. I am unable to directly approve this pull request, and recommend that others review and approve this code before merging. At a minimum, the HIGH severity comments should be addressed before merging.

mickqian · 2025-09-08T03:24:36Z

Could you check with with #8842, do we need a separate file for mtp, or we don't really need one?

yhyang201 · 2025-09-08T06:00:16Z

I think #8842 is bette, close this issue for now.

yhyang201 added 3 commits June 2, 2025 19:37

implement mtp for openbmb

5263d17

update

352aa8e

remove comment

d3fba22

yhyang201 requested review from ByronHsu, Ying1123, hnyls2002, ispobock, merrymercy, zhaochenyang20 and zhyncs as code owners June 5, 2025 15:04

gemini-code-assist Bot reviewed Jun 5, 2025

View reviewed changes

gemini-code-assist Bot suggested changes Jun 5, 2025

View reviewed changes

Comment thread python/sglang/srt/models/minicpm.py

Comment thread python/sglang/srt/configs/model_config.py

Comment thread python/sglang/srt/models/minicpm.py

Comment thread python/sglang/srt/models/minicpm_mtp.py

mickqian mentioned this pull request Jun 5, 2025

model: minicpm-4 main model #6900

Closed

6 tasks

Merge branch 'main' into openbmb

7fc91fe

yhyang201 mentioned this pull request Aug 6, 2025

feat: add eagle3 spec decoding for minicpm #8842

Open

6 tasks

yhyang201 closed this Sep 8, 2025

yhyang201 deleted the openbmb branch April 16, 2026 07:15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement MTP for MiniCPM Model#6899

Implement MTP for MiniCPM Model#6899
yhyang201 wants to merge 4 commits intosgl-project:mainfrom
yhyang201:openbmb

yhyang201 commented Jun 5, 2025 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mickqian commented Sep 8, 2025

Uh oh!

yhyang201 commented Sep 8, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

yhyang201 commented Jun 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Changelog

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Summary of Findings

Merge Readiness

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mickqian commented Sep 8, 2025

Uh oh!

yhyang201 commented Sep 8, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yhyang201 commented Jun 5, 2025 •

edited

Loading