update_dlblasanddeepep

Author: SHshenhao

lmdeploy/pytorch/backends/cuda/moe.py

@@ -474,6 +474,13 @@ def __init__(self,
         except ImportError:
             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
+        except ImportError:
+            self.use_deepep = False
+            logger.warning('For higher performance, please install DeepEP https://github.com/deepseek-ai/DeepEP')
 
         # pre-allocate buffer
         self.fusedmoe_build(True)
@@ -592,6 +599,14 @@ def _patched_fusedmoe_forward(*args, **kwargs):
 
         return deepep_moe
 
+    def update_dispatch_mode(self):
+        if self.use_deepep:
+            deepep_mode = DeepEPMode.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)
+
 
 class TritonFusedMoEBlockedF8Builder(FusedMoEBlockedF8Builder):
     """Triton fused moe blocked f8 builder."""

lmdeploy/pytorch/models/deepseek_v2.py

@@ -1158,6 +1158,11 @@ def get_input_embeddings(self):
         """Get input embeddings."""
         return self.model.get_input_embeddings()
 
+    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"):
+                self.model.layers[0].mlp.experts.impl.update_dispatch_mode()
+
     def prepare_inputs_for_generation(
         self,
         past_key_values: List[List[torch.Tensor]],
@@ -1169,11 +1174,7 @@ def prepare_inputs_for_generation(
         position_ids = context.position_ids
         attn_metadata = context.attn_metadata
 
-        from dlblas.layers.moe.token_dispatcher import DeepEPBuffer, DeepEPMode
-        deepep_mode = DeepEPMode.NORMAL
-        if context.is_decoding:
-            deepep_mode = DeepEPMode.LOW_LATENCY
-        DeepEPBuffer.set_deepep_mode(deepep_mode)
+        self._update_dispatch_mode()
 
         return dict(
             input_ids=input_ids,