qucik fix for mtp in full graph mode

Signed-off-by: chenbaixuan <[email protected]>

Author: BryanChen408

vllm_ascend/eplb/utils.py

@@ -36,7 +36,7 @@ def get_log2phy_map(self, layer_id):
         return self.layers[str(layer_id)].mtp_block.mlp.experts.get_log2phy_map()
 
 
-def get_all_expert_map(self, num_moe_layers):
+def get_all_expert_map(self, num_moe_layers=None):
     if not isinstance(self, DeepSeekMultiTokenPredictor):
         all_loads = []
         num_dense_layers = self.num_dense_layers if hasattr(

vllm_ascend/worker/model_runner_v1.py

@@ -62,7 +62,6 @@
 from vllm.model_executor.models.interfaces import (SupportsMultiModal,
                                                    supports_mrope,
                                                    supports_transcription)
-from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP
 from vllm.model_executor.models.interfaces_base import (
     VllmModelForPooling, is_pooling_model, is_text_generation_model)
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -3130,17 +3129,9 @@ def _dummy_pooler_run(
     def eplb_warmup(self):
         if self.dynamic_eplb and not self.is_eplb_warmuped:
             self.is_eplb_warmuped = True
-            mtp_instance: Optional[DeepSeekMTP] = None
-            if self.speculative_config and self.speculative_config.method == 'deepseek_mtp':
-                assert isinstance(self.drafter, MtpProposer) and isinstance(self.drafter.model, DeepSeekMTP)
-                mtp_instance=self.drafter.model
-            self.eplb_adaptor = VllmEplbAdaptor(
-                model=self.model, 
-                mtp_instance=mtp_instance, 
-                num_mtp_layers=mtp_instance.model.num_mtp_layers
-                )
+            self.eplb_adaptor = VllmEplbAdaptor(model=self.model)
             self.eplb_loader.set_adator(self.eplb_adaptor)
-            self.eplb_updator.set_adaptor(self.eplb_adaptor, mtp_instance.model.num_mtp_layers)
+            self.eplb_updator.set_adaptor(self.eplb_adaptor)
             self.eplb_updator.warm_up_eplb()
 
     def load_model(self) -> None:
@@ -3149,7 +3140,7 @@ def load_model(self) -> None:
         with DeviceMemoryProfiler() as m:  # noqa: SIM117
             self.model = get_model(vllm_config=self.vllm_config)
             if self.dynamic_eplb:
-                model_register(self.model, self.model_config)  
+                model_register(self.model, self.model_config)
             if is_310p():
                 from vllm.model_executor.layers.linear import (
                     MergedColumnParallelLinear, QKVParallelLinear,
@@ -3163,10 +3154,6 @@ def load_model(self) -> None:
             if self.drafter:
                 logger.info("Loading drafter model...")
                 self.drafter.load_model(self.model)
-                if self.speculative_config and self.speculative_config.method == 'deepseek_mtp':
-                    assert isinstance(self.drafter, MtpProposer) and isinstance(self.drafter.model, DeepSeekMTP)
-                    mtp_instance=self.drafter.model
-                model_register(mtp_instance.model, self.vllm_config)    
                 if self.drafter.name == SpecDcodeType.EAGLE3:
                     self.model.set_aux_hidden_state_layers(
                         self.model.get_eagle3_aux_hidden_state_layers())