[Misc] Avoid importing `nixl_ep` on every `vllm serve` config (#50879)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Aug 4, 2026 at 16:56 UTC
7b50d2c0bcee44ad89ee7ec75377560ed37915e3
1 file changed
+24
-23
vllm/model_executor/layers/fused_moe/all2all_utils.py
+24
-23
@@ -30,32 +30,9 @@ 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
-)
33
34
logger = init_logger(__name__)
35
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
-
36
37
def get_ep_all2all_manager(eep_stage: bool = False) -> Any:
38
if eep_stage:
@@ -93,21 +70,29 @@ def maybe_roundup_layer_hidden_size(
70
Original hidden size otherwise.
71
"""
72
if moe_parallel_config.use_deepep_ht_kernels:
73
+ from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize
74
+
75
hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
76
hidden_size, act_dtype
77
)
78
79
if moe_parallel_config.use_deepep_ll_kernels:
80
+ from .prepare_finalize.deepep_ll import DeepEPLLPrepareAndFinalize
81
+
82
hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size(
83
hidden_size
84
)
85
86
if moe_parallel_config.use_deepep_v2_kernels:
87
+ from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
88
+
89
hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size(
90
hidden_size, act_dtype
91
)
92
93
if moe_parallel_config.use_nixl_ep_kernels:
94
+ from .prepare_finalize.nixl_ep import NixlEPPrepareAndFinalize
95
+
96
hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size(
97
hidden_size
98
)
@@ -171,6 +156,8 @@ def maybe_make_prepare_finalize(
156
prepare_finalize: FusedMoEPrepareAndFinalize | None = None
157
158
if moe.use_deepep_ht_kernels:
159
+ from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize
160
+
161
assert moe.dp_size == all2all_manager.dp_world_size
162
163
all_to_all_args: dict[str, Any] = dict()
@@ -183,6 +170,11 @@ def maybe_make_prepare_finalize(
170
)
171
172
elif moe.use_deepep_ll_kernels:
173
+ from .prepare_finalize.deepep_ll import (
174
+ DEEPEP_QUANT_BLOCK_SHAPE,
175
+ DeepEPLLPrepareAndFinalize,
176
+ )
177
+
178
assert quant_config is not None
179
global_to_physical = physical_to_global = local_expert_global_ids = None
180
if routing_tables is not None:
@@ -217,6 +209,8 @@ def maybe_make_prepare_finalize(
209
local_expert_global_ids=local_expert_global_ids,
210
)
211
elif moe.use_deepep_v2_kernels:
212
+ from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
213
+
214
assert moe.dp_size == all2all_manager.dp_world_size
215
216
use_fp8_dispatch = (
@@ -247,6 +241,8 @@ def maybe_make_prepare_finalize(
241
)
242
243
elif moe.use_mori_kernels:
244
+ from .prepare_finalize.mori import MoriPrepareAndFinalize
245
+
246
assert quant_config is not None
247
248
# Note: We may want to use FP8 dispatch just to reduce
@@ -334,6 +330,11 @@ def maybe_make_prepare_finalize(
330
)
331
332
elif moe.use_nixl_ep_kernels:
333
+ from .prepare_finalize.nixl_ep import (
334
+ NIXL_EP_QUANT_BLOCK_SHAPE,
335
+ NixlEPPrepareAndFinalize,
336
+ )
337
+
338
assert quant_config is not None
339
global_to_physical = physical_to_global = local_expert_global_ids = None
340
if routing_tables is not None: