Skip to content

revert(mtp): drop shared BaseMTPModel base, keep per-model is_mtp_draft_model (revert #1337) #1339

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sufubao merged 2 commits into from
Jun 9, 2026
+111 −66
Merged

revert(mtp): drop shared BaseMTPModel base, keep per-model is_mtp_draft_model (revert #1337) #1339

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2343b3b
Revert "refactor(mtp): extract BaseMTPModel mixin shared by existing …
sufubao Jun 9, 2026
df2d921
refactor(mtp): mark MTP draft models with is_mtp_draft_model
sufubao Jun 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 0 additions & 36 deletions lightllm/models/base_mtp_model.py
View file
Open in desktop

This file was deleted.

31 changes: 29 additions & 2 deletions lightllm/models/deepseek_mtp/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,41 @@
from lightllm.models.base_mtp_model import BaseMTPModel
from typing import List
from lightllm.models.deepseek2.model import Deepseek2TpPartModel
from lightllm.models.deepseek_mtp.layer_infer.pre_layer_infer import Deepseek3MTPPreLayerInfer
from lightllm.models.deepseek_mtp.layer_weights.pre_and_post_layer_weight import Deepseek3MTPPreAndPostLayerWeight
from lightllm.common.basemodel import TpPartBaseModel


class Deepseek3MTPModel(BaseMTPModel, Deepseek2TpPartModel):
class Deepseek3MTPModel(Deepseek2TpPartModel):

# MTP draft model marker (consumed by the decode CUDA-graph / padding paths).
is_mtp_draft_model = True

pre_and_post_weight_class = Deepseek3MTPPreAndPostLayerWeight
pre_layer_infer_class = Deepseek3MTPPreLayerInfer

def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
return

def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
return

def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
return

def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
return

def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
return
Comment on lines +16 to +37

@gemini-code-assist gemini-code-assist Bot Jun 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

According to PEP 8 and standard Python idioms, explicit return statements at the end of functions that do not return a value (such as __init__ and other initializer/setter methods) are redundant and should be omitted. Removing them improves code readability and ensures consistency with Glm4MoeLiteMTPModel which correctly omits them.

Suggested change
def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
return
def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
return
def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
return
def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
return
def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
return
def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
References
  1. PEP 8 recommends omitting explicit return statements at the end of functions that do not return a value. (link)


def _init_weights(self, start_layer_index=None):
assert start_layer_index is None
self.pre_post_weight = self.pre_and_post_weight_class(
Expand Down
26 changes: 24 additions & 2 deletions lightllm/models/glm4_moe_lite_mtp/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,39 @@
from lightllm.models.base_mtp_model import BaseMTPModel
from typing import List
from lightllm.models.deepseek_mtp.layer_infer.pre_layer_infer import Deepseek3MTPPreLayerInfer
from lightllm.models.glm4_moe_lite.model import Glm4MoeLiteTpPartModel
from lightllm.models.glm4_moe_lite_mtp.layer_weights.pre_and_post_layer_weight import (
Glm4MoeLiteMTPPreAndPostLayerWeight,
)
from lightllm.common.basemodel import TpPartBaseModel
from lightllm.common.basemodel.basemodel import load_hf_weights


class Glm4MoeLiteMTPModel(BaseMTPModel, Glm4MoeLiteTpPartModel):
class Glm4MoeLiteMTPModel(Glm4MoeLiteTpPartModel):

# MTP draft model marker (consumed by the decode CUDA-graph / padding paths).
is_mtp_draft_model = True

pre_and_post_weight_class = Glm4MoeLiteMTPPreAndPostLayerWeight
pre_layer_infer_class = Deepseek3MTPPreLayerInfer

def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)

def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")

def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached

def _init_req_manager(self):
self.req_manager = self.main_model.req_manager

def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager

def _init_weights(self, start_layer_index=None):
assert start_layer_index is None

Expand Down
31 changes: 29 additions & 2 deletions lightllm/models/mistral_mtp/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,17 @@
from lightllm.models.base_mtp_model import BaseMTPModel
from typing import List
from lightllm.models.mistral.model import MistralTpPartModel
from lightllm.models.mistral_mtp.layer_weights.pre_and_post_layer_weight import MistralMTPPreAndPostLayerWeight
from lightllm.models.mistral_mtp.layer_infer.pre_layer_infer import MistralMTPPreLayerInfer
from lightllm.models.mistral_mtp.layer_infer.post_layer_infer import MistralMTPPostLayerInfer
from lightllm.models.mistral_mtp.layer_infer.transformer_layer_infer import MistralMTPTransformerLayerInfer
from lightllm.models.mistral_mtp.layer_weights.transformer_layer_weight import MistralMTPTransformerLayerWeight
from lightllm.common.basemodel import TpPartBaseModel


class MistralMTPModel(BaseMTPModel, MistralTpPartModel):
class MistralMTPModel(MistralTpPartModel):

# MTP draft model marker (consumed by the decode CUDA-graph / padding paths).
is_mtp_draft_model = True

pre_and_post_weight_class = MistralMTPPreAndPostLayerWeight
pre_layer_infer_class = MistralMTPPreLayerInfer
Expand All @@ -17,11 +21,34 @@ class MistralMTPModel(BaseMTPModel, MistralTpPartModel):

post_layer_infer_class = MistralMTPPostLayerInfer

def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
return

def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
return
Comment on lines +24 to +32

@gemini-code-assist gemini-code-assist Bot Jun 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

According to PEP 8 and standard Python idioms, explicit return statements at the end of functions that do not return a value are redundant and should be omitted. Removing them improves code readability and ensures consistency with Glm4MoeLiteMTPModel.

Suggested change
def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
return
def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
return
def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
References
  1. PEP 8 recommends omitting explicit return statements at the end of functions that do not return a value. (link)


def _init_some_value(self):
super()._init_some_value()
self.layers_num = 1
return

def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
return

def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
return

def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
return
Comment on lines +39 to +50

@gemini-code-assist gemini-code-assist Bot Jun 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

According to PEP 8 and standard Python idioms, explicit return statements at the end of functions that do not return a value are redundant and should be omitted. Removing them improves code readability and ensures consistency with Glm4MoeLiteMTPModel.

Suggested change
def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
return
def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
return
def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
return
def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
References
  1. PEP 8 recommends omitting explicit return statements at the end of functions that do not return a value. (link)


def _init_weights(self, start_layer_index=None):
assert start_layer_index is None
self.config["n_layer"] = 1
Expand Down
31 changes: 29 additions & 2 deletions lightllm/models/qwen3_moe_mtp/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,46 @@
from lightllm.models.base_mtp_model import BaseMTPModel
from typing import List
from lightllm.models.qwen3_moe.model import Qwen3MOEModel
from lightllm.models.qwen3_moe_mtp.layer_weights.pre_and_post_layer_weight import Qwen3MOEMTPPreAndPostLayerWeight
from lightllm.models.deepseek_mtp.layer_infer.pre_layer_infer import Deepseek3MTPPreLayerInfer
from lightllm.models.qwen3_moe_mtp.layer_infer.transformer_layer_infer import Qwen3MOEMTPTransformerLayerInfer
from lightllm.models.qwen3_moe_mtp.layer_weights.transformer_layer_weight import Qwen3MOEMTPTransformerLayerWeight
from lightllm.common.basemodel import TpPartBaseModel


class Qwen3MOEMTPModel(BaseMTPModel, Qwen3MOEModel):
class Qwen3MOEMTPModel(Qwen3MOEModel):

# MTP draft model marker (consumed by the decode CUDA-graph / padding paths).
is_mtp_draft_model = True

pre_and_post_weight_class = Qwen3MOEMTPPreAndPostLayerWeight
pre_layer_infer_class = Deepseek3MTPPreLayerInfer

transformer_weight_class = Qwen3MOEMTPTransformerLayerWeight
transformer_layer_infer_class = Qwen3MOEMTPTransformerLayerInfer

def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
return

def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
return

def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
return

def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
return

def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
return
Comment on lines +21 to +42

@gemini-code-assist gemini-code-assist Bot Jun 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

According to PEP 8 and standard Python idioms, explicit return statements at the end of functions that do not return a value are redundant and should be omitted. Removing them improves code readability and ensures consistency with Glm4MoeLiteMTPModel.

Suggested change
def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
return
def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
return
def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
return
def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
return
def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
return
def __init__(self, kvargs: dict):
self._pre_init(kvargs)
super().__init__(kvargs)
def _pre_init(self, kvargs: dict):
self.main_model: TpPartBaseModel = kvargs.pop("main_model")
self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
def _init_custom(self):
self._cos_cached = self.main_model._cos_cached
self._sin_cached = self.main_model._sin_cached
def _init_req_manager(self):
self.req_manager = self.main_model.req_manager
def _init_mem_manager(self):
self.mem_manager = self.main_model.mem_manager
References
  1. PEP 8 recommends omitting explicit return statements at the end of functions that do not return a value. (link)


def _init_weights(self, start_layer_index=None):
assert start_layer_index is None
mtp_index = len(self.mtp_previous_draft_models)
Expand Down
22 changes: 0 additions & 22 deletions unit_tests/models/test_base_mtp_model_mixin.py
View file
Open in desktop

This file was deleted.

Loading