A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 05:45 UTC
[MoE] Move PF Methods to Folder (#35927)
Robert Shaw committed
Mar 22, 2026 at 04:42 UTC
4383f1532e87e77b6f961e633230f47467cbd072
11 files changed
+25
-22
docs/design/moe_kernel_features.md
+4
-4
@@ -33,10 +33,10 @@ th {
33
| Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
34
| ------- | ------------------ | ------------ | ------------- | ----- | --------------------- | --------- |
35
| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
36
-| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
37
-| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
38
-| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
39
-| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize.FlashInferNVLinkOneSidedPrepareAndFinalize] |
36
+| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
37
+| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
38
+| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
39
+| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
40
41
!!! info "Table key"
42
1. All types: mxfp4, nvfp4, int4, int8, fp8
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+4
-4
@@ -199,10 +199,10 @@ register_experts(
199
200
# Disable on blackwell for now
201
if has_deep_ep() and not current_platform.has_device_capability(100):
202
- from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
202
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
203
DeepEPHTPrepareAndFinalize,
204
)
205
- from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (
205
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
206
DeepEPLLPrepareAndFinalize,
207
)
208
@@ -240,7 +240,7 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
240
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
241
FlashInferExperts,
242
)
243
- from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize import ( # noqa: E501
243
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501
244
FlashInferNVLinkTwoSidedPrepareAndFinalize,
245
)
246
@@ -271,7 +271,7 @@ if (
271
and has_flashinfer_cutlass_fused_moe()
272
and current_platform.has_device_capability(100)
273
):
274
- from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize import ( # noqa: E501
274
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided import ( # noqa: E501
275
FlashInferNVLinkOneSidedPrepareAndFinalize,
276
)
277
tests/kernels/moe/parallel_utils.py
+2
-2
@@ -19,10 +19,10 @@ from vllm.utils.import_utils import has_deep_ep
19
from vllm.utils.network_utils import get_open_port
20
21
if has_deep_ep():
22
- from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
22
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
23
DeepEPHTPrepareAndFinalize,
24
)
25
- from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (
25
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
26
DeepEPLLPrepareAndFinalize,
27
)
28
tests/kernels/moe/test_deepep_deepgemm_moe.py
+2
-2
@@ -37,10 +37,10 @@ from .parallel_utils import ProcessGroupInfo, parallel_launch
37
from .utils import make_dummy_moe_config, make_test_weights
38
39
if has_deep_ep():
40
- from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
40
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
41
DeepEPHTPrepareAndFinalize,
42
)
43
- from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (
43
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
44
DeepEPLLPrepareAndFinalize,
45
)
46
tests/kernels/moe/test_deepep_moe.py
+2
-2
@@ -32,10 +32,10 @@ from ...utils import multi_gpu_test
32
from .parallel_utils import ProcessGroupInfo, parallel_launch
33
34
if has_deep_ep():
35
- from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
35
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
36
DeepEPHTPrepareAndFinalize,
37
)
38
- from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (
38
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
39
DeepEPLLPrepareAndFinalize,
40
)
41
vllm/model_executor/layers/fused_moe/all2all_utils.py
+8
-8
@@ -15,12 +15,6 @@ from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEParallelConfig,
16
FusedMoEQuantConfig,
17
)
18
-from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize import ( # noqa: E501
19
- FlashInferNVLinkOneSidedPrepareAndFinalize,
20
-)
21
-from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize import ( # noqa: E501
22
- FlashInferNVLinkTwoSidedPrepareAndFinalize,
23
-)
18
from vllm.model_executor.layers.fused_moe.modular_kernel import (
19
FusedMoEPrepareAndFinalize,
20
)
@@ -28,6 +22,12 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
22
make_moe_prepare_and_finalize_naive_dp_ep,
23
make_moe_prepare_and_finalize_no_dp_ep,
24
)
25
+from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided import ( # noqa: E501
26
+ FlashInferNVLinkOneSidedPrepareAndFinalize,
27
+)
28
+from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501
29
+ FlashInferNVLinkTwoSidedPrepareAndFinalize,
30
+)
31
from vllm.platforms import current_platform
32
from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep
33
@@ -35,8 +35,8 @@ logger = init_logger(__name__)
35
36
if current_platform.is_cuda_alike():
37
if has_deep_ep():
38
- from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize
39
- from .deepep_ll_prepare_finalize import (
38
+ from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize
39
+ from .prepare_finalize.deepep_ll import (
40
DEEPEP_QUANT_BLOCK_SHAPE,
41
DeepEPLLPrepareAndFinalize,
42
)
vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py
+3
@@ -19,4 +19,7 @@ __all__ = [
19
"MoEPrepareAndFinalizeNoDPEPMonolithic",
20
"MoEPrepareAndFinalizeNoDPEPModular",
21
"make_moe_prepare_and_finalize_no_dp_ep",
22
+ # deepep_ht, deepep_ll, and flashinfer_a2a are not
23
+ # imported here as they have optional dependencies (deep_ep, flashinfer).
24
+ # Import them directly from their modules as needed.
25
]