From bc37fc970eb04a0e393e5517d39700acaad8aee3 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Wed, 5 Aug 2026 20:04:40 +0200 Subject: [PATCH] Revert [Misc] Avoid importing `nixl_ep` on every `vllm serve` config (#50879) (#51176) Signed-off-by: Felix Marty --- .../layers/fused_moe/all2all_utils.py | 47 +++++++++---------- 1 file changed, 23 insertions(+), 24 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 8db6fccba2..58c9c8d9f6 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -30,9 +30,32 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two FlashInferNVLinkTwoSidedPrepareAndFinalize, ) from vllm.platforms import current_platform +from vllm.utils.import_utils import ( + has_deep_ep, + has_deep_ep_v2, + has_mori, + has_nixl_ep, +) logger = init_logger(__name__) +if current_platform.is_cuda_alike(): + if has_deep_ep(): + from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize + from .prepare_finalize.deepep_ll import ( + DEEPEP_QUANT_BLOCK_SHAPE, + DeepEPLLPrepareAndFinalize, + ) + if has_deep_ep_v2(): + from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize + if has_mori(): + from .prepare_finalize.mori import MoriPrepareAndFinalize + if has_nixl_ep(): + from .prepare_finalize.nixl_ep import ( + NIXL_EP_QUANT_BLOCK_SHAPE, + NixlEPPrepareAndFinalize, + ) + def get_ep_all2all_manager(eep_stage: bool = False) -> Any: if eep_stage: @@ -70,29 +93,21 @@ def maybe_roundup_layer_hidden_size( Original hidden size otherwise. """ if moe_parallel_config.use_deepep_ht_kernels: - from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize - hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size( hidden_size, act_dtype ) if moe_parallel_config.use_deepep_ll_kernels: - from .prepare_finalize.deepep_ll import DeepEPLLPrepareAndFinalize - hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size( hidden_size ) if moe_parallel_config.use_deepep_v2_kernels: - from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize - hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size( hidden_size, act_dtype ) if moe_parallel_config.use_nixl_ep_kernels: - from .prepare_finalize.nixl_ep import NixlEPPrepareAndFinalize - hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size( hidden_size ) @@ -156,8 +171,6 @@ def maybe_make_prepare_finalize( prepare_finalize: FusedMoEPrepareAndFinalize | None = None if moe.use_deepep_ht_kernels: - from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize - assert moe.dp_size == all2all_manager.dp_world_size all_to_all_args: dict[str, Any] = dict() @@ -170,11 +183,6 @@ def maybe_make_prepare_finalize( ) elif moe.use_deepep_ll_kernels: - from .prepare_finalize.deepep_ll import ( - DEEPEP_QUANT_BLOCK_SHAPE, - DeepEPLLPrepareAndFinalize, - ) - assert quant_config is not None global_to_physical = physical_to_global = local_expert_global_ids = None if routing_tables is not None: @@ -209,8 +217,6 @@ def maybe_make_prepare_finalize( local_expert_global_ids=local_expert_global_ids, ) elif moe.use_deepep_v2_kernels: - from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize - assert moe.dp_size == all2all_manager.dp_world_size use_fp8_dispatch = ( @@ -241,8 +247,6 @@ def maybe_make_prepare_finalize( ) elif moe.use_mori_kernels: - from .prepare_finalize.mori import MoriPrepareAndFinalize - assert quant_config is not None # Note: We may want to use FP8 dispatch just to reduce @@ -330,11 +334,6 @@ def maybe_make_prepare_finalize( ) elif moe.use_nixl_ep_kernels: - from .prepare_finalize.nixl_ep import ( - NIXL_EP_QUANT_BLOCK_SHAPE, - NixlEPPrepareAndFinalize, - ) - assert quant_config is not None global_to_physical = physical_to_global = local_expert_global_ids = None if routing_tables is not None: