Moe bf16 ep (InternLM#4144)

* refactor pytorch.nn.moe

* add ep support

* fix tp

* support blocked fp8 moe with split_size<world_size

* unit test allow both fa3 and fa

* add singleton

* singleton and ctxmgrbase

* comment

* add static

* remove chunk

* remove forward dptp

* bound check

* remove monkey patch

* rename kernel

Author: grimoire

lmdeploy/pytorch/backends/cuda/graph_runner.py

@@ -5,7 +5,7 @@
 import torch
 from torch.profiler import record_function
 
-from lmdeploy.pytorch.backends.deepep_moe_checker import moe_backend
+from lmdeploy.pytorch.backends.deepep_moe_checker import get_moe_backend
 from lmdeploy.pytorch.backends.selector import get_backend
 from lmdeploy.pytorch.config import BackendConfig, CacheConfig, ModelConfig
 from lmdeploy.pytorch.model_inputs import StepContext, get_step_ctx_manager
@@ -256,7 +256,7 @@ def prepare_inputs_for_generation(
     ):
         """Prepare inputs."""
 
-        if moe_backend.use_deepep_moe_backend():
+        if get_moe_backend().use_deepep_moe_backend():
             from dlblas.layers.moe.token_dispatcher import DeepEPBuffer, DeepEPMode
             deepep_mode = DeepEPMode.LOW_LATENCY if context.is_decoding else DeepEPMode.NORMAL
             DeepEPBuffer.set_deepep_mode(deepep_mode)