Revert [Misc] Avoid importing `nixl_ep` on every `vllm serve` config (#50879) (#51176)
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
fxmarty-amd committed
Aug 5, 2026 at 20:04 UTC
bc37fc970eb04a0e393e5517d39700acaad8aee3
1 file changed
+23
-24
vllm/model_executor/layers/fused_moe/all2all_utils.py
+23
-24
@@ -30,9 +30,32 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two
30
FlashInferNVLinkTwoSidedPrepareAndFinalize,
31
)
32
from vllm.platforms import current_platform
33
+from vllm.utils.import_utils import (
34
+ has_deep_ep,
35
+ has_deep_ep_v2,
36
+ has_mori,
37
+ has_nixl_ep,
38
+)
39
40
logger = init_logger(__name__)
41
42
+if current_platform.is_cuda_alike():
43
+ if has_deep_ep():
44
+ from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize
45
+ from .prepare_finalize.deepep_ll import (
46
+ DEEPEP_QUANT_BLOCK_SHAPE,
47
+ DeepEPLLPrepareAndFinalize,
48
+ )
49
+ if has_deep_ep_v2():
50
+ from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
51
+ if has_mori():
52
+ from .prepare_finalize.mori import MoriPrepareAndFinalize
53
+ if has_nixl_ep():
54
+ from .prepare_finalize.nixl_ep import (
55
+ NIXL_EP_QUANT_BLOCK_SHAPE,
56
+ NixlEPPrepareAndFinalize,
57
+ )
58
+
59
60
def get_ep_all2all_manager(eep_stage: bool = False) -> Any:
61
if eep_stage:
@@ -70,29 +93,21 @@ def maybe_roundup_layer_hidden_size(
93
Original hidden size otherwise.
94
"""
95
if moe_parallel_config.use_deepep_ht_kernels:
73
- from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize
74
-
96
hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
97
hidden_size, act_dtype
98
)
99
100
if moe_parallel_config.use_deepep_ll_kernels:
80
- from .prepare_finalize.deepep_ll import DeepEPLLPrepareAndFinalize
81
-
101
hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size(
102
hidden_size
103
)
104
105
if moe_parallel_config.use_deepep_v2_kernels:
87
- from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
88
-
106
hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size(
107
hidden_size, act_dtype
108
)
109
110
if moe_parallel_config.use_nixl_ep_kernels:
94
- from .prepare_finalize.nixl_ep import NixlEPPrepareAndFinalize
95
-
111
hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size(
112
hidden_size
113
)
@@ -156,8 +171,6 @@ def maybe_make_prepare_finalize(
171
prepare_finalize: FusedMoEPrepareAndFinalize | None = None
172
173
if moe.use_deepep_ht_kernels:
159
- from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize
160
-
174
assert moe.dp_size == all2all_manager.dp_world_size
175
176
all_to_all_args: dict[str, Any] = dict()
@@ -170,11 +183,6 @@ def maybe_make_prepare_finalize(
183
)
184
185
elif moe.use_deepep_ll_kernels:
173
- from .prepare_finalize.deepep_ll import (
174
- DEEPEP_QUANT_BLOCK_SHAPE,
175
- DeepEPLLPrepareAndFinalize,
176
- )
177
-
186
assert quant_config is not None
187
global_to_physical = physical_to_global = local_expert_global_ids = None
188
if routing_tables is not None:
@@ -209,8 +217,6 @@ def maybe_make_prepare_finalize(
217
local_expert_global_ids=local_expert_global_ids,
218
)
219
elif moe.use_deepep_v2_kernels:
212
- from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
213
-
220
assert moe.dp_size == all2all_manager.dp_world_size
221
222
use_fp8_dispatch = (
@@ -241,8 +247,6 @@ def maybe_make_prepare_finalize(
247
)
248
249
elif moe.use_mori_kernels:
244
- from .prepare_finalize.mori import MoriPrepareAndFinalize
245
-
250
assert quant_config is not None
251
252
# Note: We may want to use FP8 dispatch just to reduce
@@ -330,11 +334,6 @@ def maybe_make_prepare_finalize(
334
)
335
336
elif moe.use_nixl_ep_kernels:
333
- from .prepare_finalize.nixl_ep import (
334
- NIXL_EP_QUANT_BLOCK_SHAPE,
335
- NixlEPPrepareAndFinalize,
336
- )
337
-
337
assert quant_config is not None
338
global_to_physical = physical_to_global = local_expert_global_ids = None
339
if routing_tables is not None: