update_dlblas_deepep

Author: SHshenhao

lmdeploy/pytorch/backends/cuda/moe.py

@@ -475,9 +475,10 @@ def __init__(self,
             self.use_deep_gemm = False
             logger.warning('For higher performance, please install DeepGEMM https://github.com/deepseek-ai/DeepGEMM')
         try:
-            import deep_ep
-            from dlblas.layers.moe.token_dispatcher import DeepEPBuffer, DeepEPMode
-            self.use_deepep = True
+            from dlblas.layers.moe.token_dispatcher import DeepEPBuffer, DeepEPMode, use_deepep
+            self.use_deepep = use_deepep
+            self.deepep_buffer = DeepEPBuffer
+            self.deepep_mode = DeepEPMode
         except ImportError:
             self.use_deepep = False
             logger.warning('For higher performance, please install DeepEP https://github.com/deepseek-ai/DeepEP')
@@ -601,11 +602,11 @@ def _patched_fusedmoe_forward(*args, **kwargs):
 
     def update_dispatch_mode(self):
         if self.use_deepep:
-            deepep_mode = DeepEPMode.NORMAL
+            deepep_mode_type = self.deepep_mode.NORMAL
             step_ctx = get_step_ctx_manager().current_context()
             if step_ctx.is_decoding:
-                deepep_mode = DeepEPMode.LOW_LATENCY
-            DeepEPBuffer.set_deepep_mode(deepep_mode)
+                deepep_mode_type = self.deepep_mode.LOW_LATENCY
+            self.deepep_buffer.set_deepep_mode(deepep_mode_type)
 
 
 class TritonFusedMoEBlockedF8Builder(FusedMoEBlockedF8Builder):

lmdeploy/pytorch/models/deepseek_v2.py

@@ -1160,7 +1160,7 @@ def get_input_embeddings(self):
 
     def _update_dispatch_mode(self):
         if isinstance(self.model.layers[0].mlp, DeepseekV2MoE):
-            if hasattr(self.model.layers[0].mlp.experts.impl, "update_dispatch_mode"):
+            if hasattr(self.model.layers[0].mlp.experts.impl, 'update_dispatch_mode'):
                 self.model.layers[0].mlp.experts.impl.update_dispatch_mode()
 
     def prepare_inputs_for_generation(