@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2754231ba3
Fork 0

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, 17:00 UTC

[Kernel] Add FlashInfer MoE A2A Kernel (#36022)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Leo Tian <lctian@nvidia.com> Co-authored-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: root <root@lyris0267.lyris.clusters.nvidia.com>

leo-cf-tian committed Mar 16, 2026 at 02:45 UTC 2754231ba3a72f41e62922d1552c33e8f3f6a9d1
19 files changed +418 -44
docs/design/moe_kernel_features.md
+2 -1
@@ -35,7 +35,8 @@ th {
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_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |
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] |
40
41 !!! info "Table key"
42 1. All types: mxfp4, nvfp4, int4, int8, fp8
docs/serving/expert_parallel_deployment.md
+2 -1
@@ -21,7 +21,8 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
21 | `allgather_reducescatter` | Default backend | Standard all2all using allgather/reducescatter primitives | General purpose, works with any EP+DP configuration |
22 | `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios |
23 | `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios |
24 -| `flashinfer_all2allv` | MNNVL systems | FlashInfer alltoallv kernels for multi-node NVLink | Systems with NVLink across nodes |
24 +| `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads |
25 +| `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
26 | `naive` | Testing/debugging | Simple broadcast-based implementation | Debugging, not recommended for production |
27
28 ## Single Node Deployment
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+38 -5
@@ -33,7 +33,10 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
33 )
34 from vllm.platforms import current_platform
35 from vllm.utils.deep_gemm import is_deep_gemm_supported
36 -from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
36 +from vllm.utils.flashinfer import (
37 + has_flashinfer_cutlass_fused_moe,
38 + has_flashinfer_nvlink_one_sided,
39 +)
40 from vllm.utils.import_utils import (
41 has_aiter,
42 has_deep_ep,
@@ -234,15 +237,15 @@ if has_mori():
237 )
238
239 if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
237 - from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501
238 - FlashInferA2APrepareAndFinalize,
239 - )
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
244 + FlashInferNVLinkTwoSidedPrepareAndFinalize,
245 + )
246
247 register_prepare_and_finalize(
245 - FlashInferA2APrepareAndFinalize,
248 + FlashInferNVLinkTwoSidedPrepareAndFinalize,
249 standard_format,
250 nvfp4_types + fp8_types,
251 blocked_quantization_support=True,
@@ -263,6 +266,36 @@ else:
266 FlashInferCutlassMoEPrepareAndFinalize = None
267 FlashInferExperts = None
268
269 +if (
270 + has_flashinfer_nvlink_one_sided()
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
275 + FlashInferNVLinkOneSidedPrepareAndFinalize,
276 + )
277 +
278 + register_prepare_and_finalize(
279 + FlashInferNVLinkOneSidedPrepareAndFinalize,
280 + standard_format,
281 + nvfp4_types,
282 + blocked_quantization_support=False,
283 + backend="flashinfer_nvlink_one_sided",
284 + supports_apply_weight_on_input=False,
285 + )
286 +
287 +if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
288 + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
289 + TrtLlmNvFp4ExpertsModular,
290 + )
291 +
292 + register_experts(
293 + TrtLlmNvFp4ExpertsModular,
294 + standard_format,
295 + nvfp4_types,
296 + blocked_quantization_support=False,
297 + supports_expert_map=True,
298 + )
299
300 if has_aiter():
301 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
vllm/config/parallel.py
+5 -2
@@ -45,7 +45,9 @@ All2AllBackend = Literal[
45 "mori",
46 "nixl_ep",
47 "allgather_reducescatter",
48 - "flashinfer_all2allv",
48 + "flashinfer_all2allv", # temporary alias for flashinfer_nvlink_two_sided
49 + "flashinfer_nvlink_two_sided",
50 + "flashinfer_nvlink_one_sided",
51 ]
52
53
@@ -158,7 +160,8 @@ class ParallelConfig:
160 - "deepep_low_latency": Use deepep low-latency kernels\n
161 - "mori": Use mori kernels\n
162 - "nixl_ep": Use nixl-ep kernels\n
161 - - "flashinfer_all2allv": Use flashinfer alltoallv kernels for mnnvl"""
163 + - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl
164 + - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels"""
165
166 max_parallel_loading_workers: int | None = None
167 """Maximum number of parallel loading workers when loading model
vllm/distributed/device_communicators/all2all.py
+132 -6
@@ -4,23 +4,36 @@ import threading
4 from typing import Any
5
6 import torch
7 +import torch.distributed as dist
8
9 import vllm.envs as envs
10 from vllm.distributed import get_dp_group, get_ep_group
11 from vllm.forward_context import get_forward_context
12 from vllm.logger import init_logger
12 -from vllm.utils.flashinfer import has_flashinfer_all2all
13 +from vllm.utils.flashinfer import (
14 + has_flashinfer_nvlink_one_sided,
15 + has_flashinfer_nvlink_two_sided,
16 +)
17 from vllm.utils.import_utils import has_deep_ep, has_mori
18
19 from .base_device_communicator import All2AllManagerBase, Cache
20
17 -if has_flashinfer_all2all():
21 +if has_flashinfer_nvlink_two_sided():
22 from flashinfer.comm import Mapping # type: ignore[import-not-found]
23 from flashinfer.comm.mnnvl import MnnvlConfig # type: ignore[import-not-found]
24 from flashinfer.comm.trtllm_alltoall import (
25 MnnvlMoe, # type: ignore[import-not-found]
26 )
27
28 +if has_flashinfer_nvlink_one_sided():
29 + from flashinfer.comm import Mapping # type: ignore[import-not-found]
30 + from flashinfer.comm.mnnvl import MnnvlConfig # type: ignore[import-not-found]
31 + from flashinfer.comm.trtllm_moe_alltoall import (
32 + MoeAlltoAll, # type: ignore[import-not-found]
33 + moe_a2a_get_workspace_size_per_rank,
34 + )
35 +
36 +
37 logger = init_logger(__name__)
38
39
@@ -529,9 +542,9 @@ class NixlEPAll2AllManager(All2AllManagerBase):
542 return 0
543
544
532 -class FlashInferAllToAllManager(All2AllManagerBase):
545 +class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
546 """
534 - All2All communication based on flashinfer kernels.
547 + All2All communication based on flashinfer all2allv/two-sided NVLink kernels.
548 """
549
550 # This type lint could be removed after all of the work in
@@ -540,7 +553,7 @@ class FlashInferAllToAllManager(All2AllManagerBase):
553 world_size: int
554
555 def __init__(self, cpu_group, tcp_store_group=None):
543 - assert has_flashinfer_all2all(), (
556 + assert has_flashinfer_nvlink_two_sided(), (
557 "flashinfer all2all module not found. Please install/check flashinfer"
558 ) # noqa
559 super().__init__(cpu_group, tcp_store_group)
@@ -597,7 +610,7 @@ class FlashInferAllToAllManager(All2AllManagerBase):
610
611 def ensure_alltoall_workspace_initialized(self):
612 """Ensure workspace is initialized"""
600 - if not has_flashinfer_all2all():
613 + if not has_flashinfer_nvlink_two_sided():
614 return False
615
616 if self.world_size <= 1:
@@ -633,6 +646,119 @@ class FlashInferAllToAllManager(All2AllManagerBase):
646 self.initialized = False
647
648
649 +class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
650 + """
651 + All2All communication based on FlashInfer's MoeAlltoAll/One-sided NVLink kernel.
652 + This is a newer kernel from trtllm that should perform better than the kernel
653 + used by flashinfer_nvlink_two_sided.
654 + """
655 +
656 + rank: int
657 + world_size: int
658 +
659 + def __init__(self, cpu_group):
660 + assert has_flashinfer_nvlink_one_sided(), (
661 + "flashinfer trtllm_moe_alltoall module not found. "
662 + "Please install/check flashinfer"
663 + )
664 + super().__init__(cpu_group)
665 + logger.debug(
666 + "Initialize FlashInfer One-sided NVLink rank=%d, world size=%d",
667 + self.rank,
668 + self.world_size,
669 + )
670 + self.initialized = False
671 + self.moe_alltoall: MoeAlltoAll | None = None
672 + self.mapping = None
673 +
674 + def initialize(
675 + self,
676 + max_num_tokens: int,
677 + top_k: int,
678 + num_experts: int,
679 + hidden_size: int,
680 + ):
681 + """Initialize the MoeAlltoAll workspace."""
682 + if self.initialized:
683 + return
684 +
685 + self.cleanup()
686 + gpus_per_node = torch.accelerator.device_count()
687 + logger.debug(
688 + "Making One-sided NVLink mapping: rank=%d, world size=%d",
689 + self.rank,
690 + self.world_size,
691 + )
692 + self.mapping = Mapping(
693 + self.world_size,
694 + self.rank,
695 + gpus_per_node,
696 + tp_size=self.world_size,
697 + moe_ep_size=self.world_size,
698 + )
699 +
700 + from vllm.distributed.device_communicators.mnnvl_compat import (
701 + CustomCommunicator,
702 + )
703 +
704 + dp_config = MnnvlConfig(
705 + comm_backend=CustomCommunicator(get_dp_group().cpu_group),
706 + )
707 + total_dispatch_payload_size_per_token = (
708 + hidden_size // 2 # nvfp4 hidden states
709 + + hidden_size // 16 # fp8 scaling factors
710 + + top_k * 4 # int32 topks ids
711 + + top_k * 4 # float32 topk weights
712 + )
713 + combine_payload_size_per_token = hidden_size * 2 # bf16 hidden states
714 + self.workspace_size = moe_a2a_get_workspace_size_per_rank(
715 + ep_size=self.world_size,
716 + max_num_tokens=max_num_tokens,
717 + total_dispatch_payload_size_per_token=total_dispatch_payload_size_per_token,
718 + combine_payload_size_per_token=combine_payload_size_per_token,
719 + )
720 +
721 + self.moe_alltoall = MoeAlltoAll(
722 + mapping=self.mapping,
723 + max_num_tokens=max_num_tokens,
724 + top_k=top_k,
725 + num_experts=num_experts,
726 + workspace_size_per_rank=self.workspace_size,
727 + mnnvl_config=dp_config,
728 + )
729 +
730 + self.gpus_per_node = gpus_per_node
731 + self.max_num_tokens = max_num_tokens
732 + self.top_k = top_k
733 + self.num_experts = num_experts
734 + self.hidden_size = hidden_size
735 + self.initialized = True
736 +
737 + logger.info(
738 + "FlashInfer One-sided NVLink initialized for rank %s, size %s",
739 + self.rank,
740 + self.world_size,
741 + )
742 + dist.barrier()
743 +
744 + def get_handle(self, kwargs):
745 + return self
746 +
747 + def cleanup(self):
748 + """Clean up resources."""
749 + if self.initialized and self.moe_alltoall is not None:
750 + try:
751 + del self.moe_alltoall
752 + except Exception as e:
753 + logger.warning(
754 + "Failed to cleanup FlashInfer One-sided NVLink workspace: %s", e
755 + )
756 + finally:
757 + self.moe_alltoall = None
758 + self.mapping = None
759 + self.initialized = False
760 +
761 +
762 class MoriAll2AllManager(All2AllManagerBase):
763 def __init__(self, cpu_group):
764 assert has_mori(), (
vllm/distributed/device_communicators/cuda_communicator.py
+17 -4
@@ -149,12 +149,25 @@ class CudaCommunicator(DeviceCommunicatorBase):
149 self.all2all_manager = NixlEPAll2AllManager(
150 self.cpu_group, tcp_store_group
151 )
152 - elif self.all2all_backend == "flashinfer_all2allv":
153 - from .all2all import FlashInferAllToAllManager
154 -
155 - self.all2all_manager = FlashInferAllToAllManager(
152 + elif (
153 + self.all2all_backend == "flashinfer_all2allv"
154 + or self.all2all_backend == "flashinfer_nvlink_two_sided"
155 + ):
156 + if self.all2all_backend == "flashinfer_all2allv":
157 + logger.warning_once(
158 + "'flashinfer_all2allv' is deprecated and has been renamed to"
159 + "'flashinfer_nvlink_two_sided'. It will be removed in a future"
160 + "release."
161 + )
162 + from .all2all import FlashInferNVLinkTwoSidedManager
163 +
164 + self.all2all_manager = FlashInferNVLinkTwoSidedManager(
165 self.cpu_group, tcp_store_group
166 )
167 + elif self.all2all_backend == "flashinfer_nvlink_one_sided":
168 + from .all2all import FlashInferNVLinkOneSidedManager
169 +
170 + self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group)
171 else:
172 raise ValueError(f"Unknown all2all backend: {self.all2all_backend}")
173
vllm/distributed/device_communicators/mnnvl_compat.py
+7 -7
@@ -5,9 +5,9 @@ from typing import Any
5 import torch.distributed as dist
6 from flashinfer.comm.mnnvl import CommBackend as CommBackend
7
8 -from vllm.utils.flashinfer import has_flashinfer_all2all
8 +from vllm.utils.flashinfer import has_flashinfer_nvlink_two_sided
9
10 -assert has_flashinfer_all2all(), "Flashinfer alltoallv module cannot be found"
10 +assert has_flashinfer_nvlink_two_sided(), "Flashinfer alltoallv module cannot be found"
11
12
13 class CustomCommunicator(CommBackend):
@@ -25,14 +25,14 @@ class CustomCommunicator(CommBackend):
25 dist.all_gather_object(gathered, data, group=self._group)
26 return gathered
27
28 - # NOTE(rob): CommBackend is an abstract class, and bcast/barrier
29 - # are unimplemented on vLLM side. If we need to utilize these
30 - # methods in the future, can create a concrete implementation.
28 def bcast(self, data: Any, root: int) -> Any:
32 - raise NotImplementedError
29 + obj_list = [data]
30 + # broadcast_object_list mutates obj_list in-place
31 + dist.broadcast_object_list(obj_list, src=root, group=self._group)
32 + return obj_list[0]
33
34 def barrier(self) -> None:
35 - raise NotImplementedError
35 + dist.barrier(group=self._group)
36
37 def Split(self, color: int, key: int) -> "CustomCommunicator":
38 return self
vllm/model_executor/layers/fused_moe/all2all_utils.py
+21 -4
@@ -5,6 +5,7 @@ from typing import Any
5
6 import torch
7
8 +from vllm.config import get_current_vllm_config
9 from vllm.distributed import (
10 get_ep_group,
11 )
@@ -14,8 +15,11 @@ from vllm.model_executor.layers.fused_moe.config import (
15 FusedMoEParallelConfig,
16 FusedMoEQuantConfig,
17 )
17 -from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import (
18 - FlashInferA2APrepareAndFinalize,
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 )
24 from vllm.model_executor.layers.fused_moe.modular_kernel import (
25 FusedMoEPrepareAndFinalize,
@@ -206,9 +210,22 @@ def maybe_make_prepare_finalize(
210 use_fp8_dispatch=use_fp8_dispatch,
211 )
212
209 - elif moe.use_fi_all2allv_kernels:
213 + elif moe.use_fi_nvl_two_sided_kernels:
214 + assert quant_config is not None
215 + prepare_finalize = FlashInferNVLinkTwoSidedPrepareAndFinalize(
216 + num_dispatchers=all2all_manager.world_size,
217 + )
218 +
219 + elif moe.use_fi_nvl_one_sided_kernels:
220 assert quant_config is not None
211 - prepare_finalize = FlashInferA2APrepareAndFinalize(
221 + max_num_tokens = (
222 + get_current_vllm_config().scheduler_config.max_num_batched_tokens
223 + )
224 + prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize(
225 + max_num_tokens=max_num_tokens,
226 + top_k=moe.experts_per_token,
227 + num_experts=moe.num_experts,
228 + hidden_size=moe.hidden_dim,
229 num_dispatchers=all2all_manager.world_size,
230 )
231
vllm/model_executor/layers/fused_moe/config.py
+16 -4
@@ -957,9 +957,17 @@ class FusedMoEParallelConfig:
957 return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency"
958
959 @property
960 - def use_fi_all2allv_kernels(self):
960 + def use_fi_nvl_two_sided_kernels(self):
961 + return self.use_all2all_kernels and (
962 + self.all2all_backend == "flashinfer_all2allv"
963 + or self.all2all_backend == "flashinfer_nvlink_two_sided"
964 + )
965 +
966 + @property
967 + def use_fi_nvl_one_sided_kernels(self):
968 return (
962 - self.use_all2all_kernels and self.all2all_backend == "flashinfer_all2allv"
969 + self.use_all2all_kernels
970 + and self.all2all_backend == "flashinfer_nvlink_one_sided"
971 )
972
973 @property
@@ -1240,8 +1248,12 @@ class FusedMoEConfig:
1248 return self.moe_parallel_config.use_mori_kernels
1249
1250 @property
1243 - def use_fi_all2allv_kernels(self):
1244 - return self.moe_parallel_config.use_fi_all2allv_kernels
1251 + def use_fi_nvl_two_sided_kernels(self):
1252 + return self.moe_parallel_config.use_fi_nvl_two_sided_kernels
1253 +
1254 + @property
1255 + def use_fi_nvl_one_sided_kernels(self):
1256 + return self.moe_parallel_config.use_fi_nvl_one_sided_kernels
1257
1258 @property
1259 def use_naive_all2all_kernels(self):
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+2 -1
@@ -396,8 +396,9 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):
396 # Note that the BATCHED activation format does not use
397 # the expert map for identifying experts.
398 return not (
399 - moe_parallel_config.use_fi_all2allv_kernels
399 + moe_parallel_config.use_fi_nvl_two_sided_kernels
400 or moe_parallel_config.use_deepep_ht_kernels
401 + or moe_parallel_config.use_fi_nvl_one_sided_kernels
402 )
403
404 def supports_expert_map(self) -> bool:
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+4 -1
@@ -152,7 +152,10 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
152 @staticmethod
153 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
154 # NOTE(rob): discovered an IMA with this combination. Needs investigation.
155 - return not moe_parallel_config.use_fi_all2allv_kernels
155 + return not (
156 + moe_parallel_config.use_fi_nvl_two_sided_kernels
157 + or moe_parallel_config.use_fi_nvl_one_sided_kernels
158 + )
159
160 def supports_expert_map(self) -> bool:
161 return True
vllm/model_executor/layers/fused_moe/flashinfer_nvlink_one_sided_prepare_finalize.py new
+146
@@ -0,0 +1,146 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
6 +from vllm.distributed import get_ep_group
7 +from vllm.forward_context import get_forward_context
8 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9 +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
10 +from vllm.utils.flashinfer import nvfp4_block_scale_interleave
11 +
12 +
13 +def get_local_sizes():
14 + return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
15 +
16 +
17 +class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
18 + """FlashInfer implementation using the Moe AlltoAll kernel."""
19 +
20 + def __init__(
21 + self,
22 + max_num_tokens: int,
23 + top_k: int,
24 + num_experts: int,
25 + hidden_size: int,
26 + num_dispatchers: int = 1,
27 + ):
28 + super().__init__()
29 + self.max_num_tokens = max_num_tokens
30 + self.top_k = top_k
31 + self.num_experts = num_experts
32 + self.hidden_size = hidden_size
33 + self.num_dispatchers_ = num_dispatchers
34 +
35 + self.all2all_manager = get_ep_group().device_communicator.all2all_manager
36 + self.all2all_manager.initialize(
37 + max_num_tokens=self.max_num_tokens,
38 + top_k=self.top_k,
39 + num_experts=self.num_experts,
40 + hidden_size=self.hidden_size,
41 + )
42 +
43 + @property
44 + def activation_format(self) -> mk.FusedMoEActivationFormat:
45 + return mk.FusedMoEActivationFormat.Standard
46 +
47 + def max_num_tokens_per_rank(self) -> int | None:
48 + return None
49 +
50 + def num_dispatchers(self) -> int:
51 + return self.num_dispatchers_
52 +
53 + def output_is_reduced(self) -> bool:
54 + return False
55 +
56 + def topk_indices_dtype(self) -> torch.dtype | None:
57 + return torch.int32
58 +
59 + def prepare(
60 + self,
61 + a1: torch.Tensor,
62 + topk_weights: torch.Tensor,
63 + topk_ids: torch.Tensor,
64 + num_experts: int,
65 + expert_map: torch.Tensor | None,
66 + apply_router_weight_on_input: bool,
67 + quant_config: FusedMoEQuantConfig,
68 + defer_input_quant: bool = False,
69 + ) -> mk.PrepareResultType:
70 + if apply_router_weight_on_input:
71 + topk = topk_ids.size(1)
72 + assert topk == 1, (
73 + "apply_router_weight_on_input is only implemented for topk=1"
74 + )
75 + a1.mul_(topk_weights.to(a1.dtype))
76 +
77 + global_num_tokens_cpu = get_local_sizes()
78 + self.runtime_max_tokens_per_rank = (
79 + max(global_num_tokens_cpu)
80 + if global_num_tokens_cpu is not None
81 + else a1.shape[0]
82 + )
83 +
84 + a1q, a1q_scale = moe_kernel_quantize_input(
85 + a1,
86 + quant_config.a1_gscale,
87 + quant_config.quant_dtype,
88 + quant_config.per_act_token_quant,
89 + quant_config.block_shape,
90 + is_fp4_scale_swizzled=False, # delay swizzle to after comm
91 + )
92 +
93 + payloads = []
94 + payloads.append(a1q)
95 + if a1q_scale is not None:
96 + payloads.append(a1q_scale)
97 + payloads.append(topk_ids)
98 + payloads.append(topk_weights)
99 +
100 + recv_payloads = self.all2all_manager.moe_alltoall.dispatch(
101 + token_selected_experts=topk_ids,
102 + input_payloads=payloads,
103 + runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
104 + )
105 + if a1q_scale is not None:
106 + a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
107 + # Apply scale interleaving only for CUTLASS (not TRT-LLM)
108 + if (
109 + quant_config.quant_dtype == "nvfp4"
110 + and quant_config.is_nvfp4_scale_swizzled
111 + ):
112 + a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])
113 + a1q_scale_recv = a1q_scale_recv.view(torch.uint8)
114 + a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)
115 + a1q_scale_recv = a1q_scale_recv.view(-1, self.hidden_size // 16)
116 + else:
117 + a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads
118 + a1q_scale_recv = None
119 + a1q_recv = a1q_recv.view(-1, a1q_recv.shape[-1])
120 + topk_ids_recv = topk_ids_recv.view(-1, topk_ids_recv.shape[-1])
121 + topk_weights_recv = topk_weights_recv.view(-1, topk_weights_recv.shape[-1])
122 +
123 + return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv
124 +
125 + def finalize(
126 + self,
127 + output: torch.Tensor,
128 + fused_expert_output: torch.Tensor,
129 + topk_weights: torch.Tensor,
130 + topk_ids: torch.Tensor,
131 + apply_router_weight_on_input: bool,
132 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
133 + ) -> None:
134 + assert self.all2all_manager.moe_alltoall is not None
135 +
136 + ep_size = self.all2all_manager.world_size
137 + hidden_size = fused_expert_output.shape[-1]
138 + fused_expert_output = fused_expert_output.view(
139 + ep_size, self.runtime_max_tokens_per_rank, hidden_size
140 + )
141 +
142 + combined_output = self.all2all_manager.moe_alltoall.combine(
143 + payload=fused_expert_output,
144 + runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
145 + )
146 + output.copy_(combined_output)
vllm/model_executor/layers/fused_moe/flashinfer_nvlink_two_sided_prepare_finalize.py renamed
+1 -1
@@ -18,7 +18,7 @@ def get_local_sizes():
18 return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
19
20
21 -class FlashInferA2APrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
21 +class FlashInferNVLinkTwoSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
22 """Base class for FlashInfer MoE prepare and finalize operations."""
23
24 def __init__(
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+4 -1
@@ -600,7 +600,10 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
600
601 @staticmethod
602 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
603 - return not moe_parallel_config.use_fi_all2allv_kernels
603 + return not (
604 + moe_parallel_config.use_fi_nvl_two_sided_kernels
605 + or moe_parallel_config.use_fi_nvl_one_sided_kernels
606 + )
607
608 @property
609 def quant_type_id(self) -> int:
vllm/model_executor/layers/fused_moe/fused_moe.py
+4 -1
@@ -1965,7 +1965,10 @@ class TritonExperts(mk.FusedMoEExpertsModular):
1965
1966 @staticmethod
1967 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
1968 - return not moe_parallel_config.use_fi_all2allv_kernels
1968 + return not (
1969 + moe_parallel_config.use_fi_nvl_two_sided_kernels
1970 + or moe_parallel_config.use_fi_nvl_one_sided_kernels
1971 + )
1972
1973 def supports_expert_map(self) -> bool:
1974 return True
vllm/model_executor/layers/fused_moe/layer.py
+1 -1
@@ -638,7 +638,7 @@ class FusedMoE(CustomOp):
638 self.use_overlapped = (
639 not (
640 (self.enable_eplb and backend != "allgather_reducescatter")
641 - or self.moe_parallel_config.use_fi_all2allv_kernels
641 + or self.moe_parallel_config.use_fi_nvl_two_sided_kernels
642 )
643 and self._shared_experts is not None
644 )
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+4 -1
@@ -332,7 +332,10 @@ class AiterExperts(mk.FusedMoEExpertsModular):
332
333 @staticmethod
334 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
335 - return not moe_parallel_config.use_fi_all2allv_kernels
335 + return not (
336 + moe_parallel_config.use_fi_nvl_two_sided_kernels
337 + or moe_parallel_config.use_fi_nvl_one_sided_kernels
338 + )
339
340 def supports_expert_map(self):
341 return True
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+1 -1
@@ -233,7 +233,7 @@ class DefaultMoERunner(MoERunner):
233 return (
234 self.moe_config.moe_parallel_config.use_deepep_ll_kernels
235 or self.moe_config.moe_parallel_config.use_mori_kernels
236 - or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels
236 + or self.moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels
237 or self.moe_config.moe_parallel_config.use_nixl_ep_kernels
238 ) and envs.VLLM_ENABLE_MOE_DP_CHUNK
239
vllm/utils/flashinfer.py
+11 -2
@@ -150,7 +150,7 @@ def has_flashinfer_comm() -> bool:
150
151
152 @functools.cache
153 -def has_flashinfer_all2all() -> bool:
153 +def has_flashinfer_nvlink_two_sided() -> bool:
154 """Return `True` if FlashInfer mnnvl all2all is available."""
155 if not has_flashinfer_comm():
156 return False
@@ -170,6 +170,14 @@ def has_flashinfer_all2all() -> bool:
170 return True
171
172
173 +@functools.cache
174 +def has_flashinfer_nvlink_one_sided() -> bool:
175 + """Return `True` if FlashInfer trtllm_moe_alltoall module is available."""
176 + if not has_flashinfer_comm():
177 + return False
178 + return importlib.util.find_spec("flashinfer.comm.trtllm_moe_alltoall") is not None
179 +
180 +
181 @functools.cache
182 def has_flashinfer_moe() -> bool:
183 """Return `True` if FlashInfer MoE module is available."""
@@ -766,7 +774,8 @@ __all__ = [
774 "autotune",
775 "has_flashinfer_moe",
776 "has_flashinfer_comm",
769 - "has_flashinfer_all2all",
777 + "has_flashinfer_nvlink_two_sided",
778 + "has_flashinfer_nvlink_one_sided",
779 "has_flashinfer_cutlass_fused_moe",
780 "has_flashinfer_cutedsl_grouped_gemm_nt_masked",
781 "has_flashinfer_fp8_blockscale_gemm",