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

A high-throughput and memory-efficient inference and serving engine for LLMs

[WideEP] Integrate DeepEP v2 (#41183)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Tyler Michael Smith committed Jun 8, 2026 at 21:07 UTC e2f993dc4116ba96ea10ec40b7bd2ee45b27b2b8
22 files changed +1409 -17
.buildkite/test_areas/kernels.yaml
+1
@@ -299,3 +299,4 @@ steps:
299 - vllm/config
300 commands:
301 - pytest -v -s kernels/moe/test_moe_layer.py
302 + - pytest -v -s kernels/moe/test_deepep_v2_moe.py
tests/distributed/test_mnnvl_alltoall.py
+78
@@ -19,6 +19,7 @@ from vllm.utils.flashinfer import (
19 has_flashinfer_nvlink_one_sided,
20 has_flashinfer_nvlink_two_sided,
21 )
22 +from vllm.utils.import_utils import has_deep_ep_v2
23 from vllm.utils.network_utils import get_open_port
24
25 from ..utils import init_test_distributed_environment
@@ -194,6 +195,10 @@ requires_ptrace = pytest.mark.skipif(
195 not _has_sys_ptrace(),
196 reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)",
197 )
198 +requires_deep_ep_v2 = pytest.mark.skipif(
199 + not has_deep_ep_v2(),
200 + reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4",
201 +)
202
203 # NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have
204 # their own @requires_two_sided / @requires_one_sided decorators, and
@@ -856,3 +861,76 @@ def _one_sided_data_worker(rank, world_size):
861 def test_one_sided_dispatch_combine(world_size):
862 """Test FlashInfer one-sided dispatch/combine with actual data flow."""
863 _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)
864 +
865 +
866 +# ---------------------------------------------------------------------------
867 +# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle
868 +# ---------------------------------------------------------------------------
869 +#
870 +# Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using
871 +# the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4.
872 +#
873 +# Uses EP group because the DeepEP v2 manager is constructed with an
874 +# EP-scoped communicator in production. With tp=world_size the EP group
875 +# spans all ranks.
876 +# ---------------------------------------------------------------------------
877 +
878 +
879 +def _deepep_v2_lifecycle_worker(rank, world_size):
880 + from vllm.distributed.device_communicators.all2all import (
881 + DeepEPV2All2AllManager,
882 + )
883 +
884 + cpu_group = get_ep_group().cpu_group
885 + manager = DeepEPV2All2AllManager(cpu_group)
886 +
887 + assert manager.rank == rank
888 + assert manager.world_size == world_size
889 + assert manager._num_sms is None
890 +
891 + hidden_size = 7168
892 + num_experts = world_size * 32
893 + num_topk = 8
894 + max_tokens = 256
895 +
896 + handle_kwargs = dict(
897 + num_max_tokens_per_rank=max_tokens,
898 + hidden=hidden_size,
899 + num_topk=num_topk,
900 + num_experts=num_experts,
901 + use_fp8_dispatch=False,
902 + )
903 +
904 + handle = manager.get_handle(handle_kwargs)
905 + assert handle is not None
906 + assert manager._num_sms is not None
907 + assert manager._num_sms > 0
908 +
909 + torch.distributed.barrier()
910 +
911 + # get_handle again with same args should return cached handle
912 + handle2 = manager.get_handle(dict(handle_kwargs))
913 + assert handle2 is handle
914 +
915 + torch.distributed.barrier()
916 +
917 + # Destroy clears the cache
918 + manager.destroy()
919 + assert len(manager.handle_cache._cache) == 0
920 +
921 + torch.distributed.barrier()
922 +
923 + # Re-create after destroy
924 + handle3 = manager.get_handle(dict(handle_kwargs))
925 + assert handle3 is not None
926 +
927 + torch.distributed.barrier()
928 + manager.destroy()
929 +
930 +
931 +@requires_multi_gpu
932 +@requires_deep_ep_v2
933 +@pytest.mark.parametrize("world_size", [2])
934 +def test_deepep_v2_manager_lifecycle(world_size):
935 + """Test DeepEP v2 ElasticBuffer manager init, caching, and destroy."""
936 + _spawn_workers(_deepep_v2_lifecycle_worker, world_size)
tests/kernels/moe/modular_kernel_tools/common.py
+67
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4 +from types import SimpleNamespace
5 from typing import Any
6
7 import torch
@@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
44 from vllm.utils.import_utils import (
45 has_aiter,
46 has_deep_ep,
47 + has_deep_ep_v2,
48 has_deep_gemm,
49 has_mori,
50 )
@@ -150,6 +152,10 @@ class Config:
152 make env data for vllm launch.
153 """
154 vllm_config = VllmConfig()
155 + vllm_config.model_config = SimpleNamespace(
156 + enforce_eager=True,
157 + is_moe=True,
158 + )
159 vllm_config.parallel_config.data_parallel_size = self.world_size
160 vllm_config.parallel_config.enable_expert_parallel = True
161
@@ -239,6 +245,10 @@ class Config:
245 or info.backend == "deepep_low_latency"
246 )
247
248 + def needs_deep_ep_v2(self):
249 + info = prepare_finalize_info(self.prepare_finalize_type)
250 + return info.backend == "deepep_v2"
251 +
252 def needs_aiter(self):
253 info = expert_info(self.fused_experts_type)
254 return info.needs_aiter
@@ -315,6 +325,8 @@ class Config:
325 # Check dependencies (turn into asserts?)
326 if self.needs_deep_ep() and not has_deep_ep():
327 return False, "Needs DeepEP, but DeepEP not available."
328 + if self.needs_deep_ep_v2() and not has_deep_ep_v2():
329 + return False, "Needs DeepEP v2, but DeepEP v2 not available."
330 if self.needs_deep_gemm() and not has_deep_gemm():
331 return False, "Needs DeepGEMM, but DeepGEMM not available."
332 if self.needs_aiter() and not has_aiter(): # noqa: SIM103
@@ -657,6 +669,58 @@ def make_modular_kernel(
669 return modular_kernel
670
671
672 +def _maybe_convert_weights_for_experts(
673 + config: Config,
674 + rank_weights: WeightTensors,
675 +) -> WeightTensors:
676 + """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK)."""
677 + from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
678 + Fp8MoeBackend,
679 + convert_to_fp8_moe_kernel_format,
680 + )
681 +
682 + fe_type = config.fused_experts_type
683 + fe_name = getattr(fe_type, "__name__", "")
684 +
685 + backend: Fp8MoeBackend | None = None
686 + if fe_name == "TrtLlmFp8ExpertsModular":
687 + backend = Fp8MoeBackend.FLASHINFER_TRTLLM
688 + elif fe_name == "FlashInferExperts":
689 + backend = Fp8MoeBackend.FLASHINFER_CUTLASS
690 +
691 + if backend is None or not rank_weights.is_quantized():
692 + return rank_weights
693 +
694 + mock_layer = SimpleNamespace(
695 + weight_block_size=config.quant_block_shape,
696 + moe_config=SimpleNamespace(
697 + is_act_and_mul=True,
698 + intermediate_size_per_partition=config.N,
699 + ),
700 + activation=SimpleNamespace(is_gated=True),
701 + )
702 +
703 + w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format(
704 + fp8_backend=backend,
705 + layer=mock_layer,
706 + w13=rank_weights.w1,
707 + w2=rank_weights.w2,
708 + w13_scale=rank_weights.w1_scale,
709 + w2_scale=rank_weights.w2_scale,
710 + w13_input_scale=None,
711 + w2_input_scale=None,
712 + )
713 +
714 + return WeightTensors(
715 + w1=w1,
716 + w2=w2,
717 + w1_scale=w1_scale,
718 + w2_scale=w2_scale,
719 + w1_gs=rank_weights.w1_gs,
720 + w2_gs=rank_weights.w2_gs,
721 + )
722 +
723 +
724 def run_modular_kernel(
725 pgi: ProcessGroupInfo,
726 vllm_config: VllmConfig,
@@ -669,6 +733,7 @@ def run_modular_kernel(
733
734 # weights for rank
735 rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
736 + rank_weights = _maybe_convert_weights_for_experts(config, rank_weights)
737
738 if config.quant_dtype == "nvfp4":
739 gscale = _make_gscale(config.num_local_experts)
@@ -714,6 +779,8 @@ def run_modular_kernel(
779 [num_tokens] * config.world_size, device="cuda", dtype=torch.int
780 )
781
782 + torch.distributed.barrier()
783 +
784 with set_forward_context(
785 None,
786 vllm_config,
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+27
@@ -36,10 +36,12 @@ from vllm.utils.deep_gemm import is_deep_gemm_supported
36 from vllm.utils.flashinfer import (
37 has_flashinfer_cutlass_fused_moe,
38 has_flashinfer_nvlink_one_sided,
39 + has_flashinfer_trtllm_fused_moe,
40 )
41 from vllm.utils.import_utils import (
42 has_aiter,
43 has_deep_ep,
44 + has_deep_ep_v2,
45 has_deep_gemm,
46 has_mori,
47 )
@@ -216,6 +218,19 @@ if has_deep_ep() and not current_platform.has_device_capability(100):
218 backend="deepep_low_latency",
219 )
220
221 +if has_deep_ep_v2() and current_platform.has_device_capability(100):
222 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
223 + DeepEPV2PrepareAndFinalize,
224 + )
225 +
226 + register_prepare_and_finalize(
227 + DeepEPV2PrepareAndFinalize,
228 + standard_format,
229 + common_float_types,
230 + blocked_quantization_support=True,
231 + backend="deepep_v2",
232 + )
233 +
234 if has_mori():
235 from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
236 MoriPrepareAndFinalize,
@@ -289,6 +304,18 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
304 blocked_quantization_support=False,
305 )
306
307 +if has_flashinfer_trtllm_fused_moe() and current_platform.has_device_capability(100):
308 + from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
309 + TrtLlmFp8ExpertsModular,
310 + )
311 +
312 + register_experts(
313 + TrtLlmFp8ExpertsModular,
314 + standard_format,
315 + fp8_types,
316 + blocked_quantization_support=True,
317 + )
318 +
319 if has_aiter():
320 from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
321 AiterExperts,
tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+8 -1
@@ -106,7 +106,7 @@ def _worker_parallel_launch(
106 if vllm_config is not None:
107 cpu_group = _set_vllm_config(vllm_config, world_size, rank, local_rank)
108
109 - try:
109 + def _run_worker():
110 worker(
111 ProcessGroupInfo(
112 world_size=world_size,
@@ -121,6 +121,13 @@ def _worker_parallel_launch(
121 *args,
122 **worker_kwargs,
123 )
124 +
125 + try:
126 + if vllm_config is not None:
127 + with set_current_vllm_config(vllm_config):
128 + _run_worker()
129 + else:
130 + _run_worker()
131 except Exception as ex:
132 print(ex)
133 traceback.print_exc()
tests/kernels/moe/parallel_utils.py
+46 -3
@@ -15,7 +15,7 @@ from torch.distributed import ProcessGroup
15 from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]
16 from typing_extensions import ParamSpec
17
18 -from vllm.utils.import_utils import has_deep_ep
18 +from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2
19 from vllm.utils.network_utils import get_open_port
20
21 if has_deep_ep():
@@ -26,6 +26,11 @@ if has_deep_ep():
26 DeepEPLLPrepareAndFinalize,
27 )
28
29 +if has_deep_ep_v2():
30 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
31 + DeepEPV2PrepareAndFinalize,
32 + )
33 +
34 ## Parallel Processes Utils
35
36 P = ParamSpec("P")
@@ -55,11 +60,10 @@ def _worker_parallel_launch(
60 torch.accelerator.set_device_index(local_rank)
61 device = torch.device("cuda", local_rank)
62 torch.distributed.init_process_group(
58 - backend="cpu:gloo,cuda:nccl",
63 + backend="nccl",
64 init_method=init_method,
65 rank=rank,
66 world_size=world_size,
62 - device_id=device,
67 )
68 barrier = torch.tensor([rank], device=device)
69 torch.distributed.all_reduce(barrier)
@@ -200,3 +204,42 @@ def make_deepep_a2a(
204
205 assert deepep_ll_args is not None
206 return make_deepep_ll_a2a(pg, pgi, deepep_ll_args, q_dtype, block_shape)
207 +
208 +
209 +@dataclasses.dataclass
210 +class DeepEPV2Args:
211 + num_local_experts: int
212 + num_experts: int
213 + num_topk: int
214 + hidden_size: int
215 + max_tokens_per_rank: int
216 + use_fp8_dispatch: bool
217 +
218 +
219 +def make_deepep_v2_a2a(
220 + pg: ProcessGroup,
221 + pgi: ProcessGroupInfo,
222 + dp_size: int,
223 + v2_args: DeepEPV2Args,
224 + use_cudagraph: bool = False,
225 +):
226 + import deep_ep
227 +
228 + buffer = deep_ep.ElasticBuffer(
229 + group=pg,
230 + num_max_tokens_per_rank=v2_args.max_tokens_per_rank,
231 + hidden=v2_args.hidden_size,
232 + num_topk=v2_args.num_topk,
233 + use_fp8_dispatch=v2_args.use_fp8_dispatch,
234 + explicitly_destroy=True,
235 + )
236 + return DeepEPV2PrepareAndFinalize(
237 + buffer=buffer,
238 + num_dispatchers=pgi.world_size,
239 + dp_size=dp_size,
240 + rank_expert_offset=pgi.rank * v2_args.num_local_experts,
241 + num_experts=v2_args.num_experts,
242 + num_topk=v2_args.num_topk,
243 + use_fp8_dispatch=v2_args.use_fp8_dispatch,
244 + use_cudagraph=use_cudagraph,
245 + )
tests/kernels/moe/test_deepep_v2_moe.py new
+542
@@ -0,0 +1,542 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Test DeepEP v2 (ElasticBuffer) dispatch-combine logic.
5 +Compares against a pure-PyTorch reference MoE implementation.
6 +"""
7 +
8 +import dataclasses
9 +
10 +import pytest
11 +import torch.distributed
12 +from torch.distributed import ProcessGroup
13 +
14 +from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights
15 +from tests.kernels.utils import torch_experts
16 +from vllm.config import VllmConfig, set_current_vllm_config
17 +from vllm.model_executor.layers.fused_moe import TritonExperts
18 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19 +from vllm.model_executor.layers.fused_moe.config import (
20 + FusedMoEQuantConfig,
21 +)
22 +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
23 +from vllm.utils.import_utils import has_deep_ep_v2
24 +from vllm.utils.torch_utils import set_random_seed
25 +from vllm.v1.worker.workspace import init_workspace_manager
26 +
27 +from ...utils import multi_gpu_test
28 +from .parallel_utils import ProcessGroupInfo, parallel_launch
29 +
30 +if has_deep_ep_v2():
31 + from .parallel_utils import DeepEPV2Args, make_deepep_v2_a2a
32 +
33 +requires_deep_ep_v2 = pytest.mark.skipif(
34 + not has_deep_ep_v2(),
35 + reason="Requires DeepEP v2 (ElasticBuffer)",
36 +)
37 +
38 +
39 +@dataclasses.dataclass
40 +class TestConfig:
41 + dtype: torch.dtype
42 + topk: int
43 + m: int
44 + k: int
45 + n: int
46 + num_experts: int
47 +
48 +
49 +@dataclasses.dataclass
50 +class TestTensors:
51 + rank_tokens: torch.Tensor
52 + rank_token_scales: torch.Tensor | None
53 + topk: torch.Tensor
54 + topk_weights: torch.Tensor
55 + config: TestConfig
56 +
57 + @staticmethod
58 + def make(config: TestConfig) -> "TestTensors":
59 + assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn]
60 + token_dtype = (
61 + torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype
62 + )
63 + rank_tokens = (
64 + torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10
65 + )
66 +
67 + topk = torch.stack(
68 + [
69 + torch.randperm(config.num_experts, device="cuda")[: config.topk]
70 + for _ in range(config.m)
71 + ]
72 + ).to(dtype=torch.int64)
73 + topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda")
74 + return TestTensors(
75 + rank_tokens=rank_tokens,
76 + rank_token_scales=None,
77 + topk=topk,
78 + topk_weights=topk_weights,
79 + config=config,
80 + )
81 +
82 +
83 +def make_modular_kernel(
84 + pg: ProcessGroup,
85 + pgi: ProcessGroupInfo,
86 + dp_size: int,
87 + hidden_size: int,
88 + num_experts: int,
89 + num_local_experts: int,
90 + topk: int,
91 + q_dtype: torch.dtype | None,
92 + use_fp8_dispatch: bool,
93 + quant_config: FusedMoEQuantConfig,
94 + use_cudagraph: bool = False,
95 +) -> FusedMoEKernel:
96 + v2_args = DeepEPV2Args(
97 + num_local_experts=num_local_experts,
98 + num_experts=num_experts,
99 + num_topk=topk,
100 + hidden_size=hidden_size,
101 + max_tokens_per_rank=8192,
102 + use_fp8_dispatch=use_fp8_dispatch,
103 + )
104 +
105 + a2a = make_deepep_v2_a2a(
106 + pg=pg,
107 + pgi=pgi,
108 + dp_size=dp_size,
109 + v2_args=v2_args,
110 + use_cudagraph=use_cudagraph,
111 + )
112 +
113 + moe_config = make_dummy_moe_config(
114 + num_experts=num_local_experts,
115 + experts_per_token=topk,
116 + hidden_dim=hidden_size,
117 + )
118 +
119 + fused_experts = TritonExperts(
120 + moe_config=moe_config,
121 + quant_config=quant_config,
122 + )
123 +
124 + mk = FusedMoEKernel(
125 + prepare_finalize=a2a,
126 + fused_experts=fused_experts,
127 + inplace=False,
128 + )
129 + return mk
130 +
131 +
132 +def deepep_v2_moe_impl(
133 + pg: ProcessGroup,
134 + pgi: ProcessGroupInfo,
135 + dp_size: int,
136 + test_tensors: TestTensors,
137 + w1: torch.Tensor,
138 + w2: torch.Tensor,
139 + w1_scale: torch.Tensor | None,
140 + w2_scale: torch.Tensor | None,
141 + num_experts: int,
142 + topk: int,
143 + use_fp8_dispatch: bool,
144 + per_act_token_quant: bool,
145 +) -> torch.Tensor:
146 + num_local_experts = w1.size(0)
147 +
148 + def build_expert_map():
149 + expert_map = torch.full((num_experts,), fill_value=-1, dtype=torch.int32)
150 + s = pgi.rank * num_local_experts
151 + e = s + num_local_experts
152 + expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
153 + device = torch.accelerator.current_device_index()
154 + return expert_map.to(device=device, dtype=torch.int32)
155 +
156 + is_quantized = w1.dtype == torch.float8_e4m3fn
157 + q_dtype = torch.float8_e4m3fn if is_quantized else None
158 +
159 + quant_config = FusedMoEQuantConfig.make(
160 + q_dtype,
161 + w1_scale=w1_scale,
162 + w2_scale=w2_scale,
163 + per_act_token_quant=per_act_token_quant,
164 + a1_scale=test_tensors.rank_token_scales,
165 + )
166 +
167 + hidden_size = test_tensors.rank_tokens.size(1)
168 +
169 + mk: FusedMoEKernel = make_modular_kernel(
170 + pg,
171 + pgi,
172 + dp_size,
173 + hidden_size,
174 + num_experts,
175 + num_local_experts,
176 + topk,
177 + q_dtype,
178 + use_fp8_dispatch,
179 + quant_config,
180 + )
181 +
182 + out = mk.apply(
183 + hidden_states=test_tensors.rank_tokens,
184 + w1=w1,
185 + w2=w2,
186 + topk_weights=test_tensors.topk_weights,
187 + topk_ids=test_tensors.topk,
188 + activation=MoEActivation.SILU,
189 + global_num_experts=num_experts,
190 + expert_map=build_expert_map(),
191 + apply_router_weight_on_input=False,
192 + )
193 +
194 + return out
195 +
196 +
197 +def _deep_ep_v2_moe(
198 + pgi: ProcessGroupInfo,
199 + dp_size: int,
200 + config: TestConfig,
201 + w1: torch.Tensor,
202 + w2: torch.Tensor,
203 + w1_scale: torch.Tensor | None,
204 + w2_scale: torch.Tensor | None,
205 + use_fp8_dispatch: bool,
206 + per_act_token_quant: bool,
207 +):
208 + device = torch.device(f"cuda:{pgi.local_rank}")
209 + init_workspace_manager(device)
210 +
211 + is_quantized = w1.dtype == torch.float8_e4m3fn
212 + device_idx = torch.accelerator.current_device_index()
213 + w1 = w1.to(device=device_idx)
214 + w2 = w2.to(device=device_idx)
215 + if is_quantized:
216 + assert w1_scale is not None and w2_scale is not None
217 + w1_scale = w1_scale.to(device=device_idx)
218 + w2_scale = w2_scale.to(device=device_idx)
219 +
220 + pg = torch.distributed.new_group(list(range(pgi.world_size)))
221 + test_tensors = TestTensors.make(config)
222 +
223 + with set_current_vllm_config(VllmConfig()):
224 + # Reference
225 + q_dtype = torch.float8_e4m3fn if is_quantized else None
226 + torch_combined = torch_experts(
227 + test_tensors.rank_tokens,
228 + w1,
229 + w2,
230 + test_tensors.topk_weights,
231 + test_tensors.topk,
232 + w1_scale=w1_scale,
233 + w2_scale=w2_scale,
234 + quant_dtype=q_dtype,
235 + per_act_token_quant=per_act_token_quant,
236 + )
237 +
238 + # Splice experts for this rank
239 + num_local_experts = config.num_experts // pgi.world_size
240 + e_start = num_local_experts * pgi.rank
241 + e_end = e_start + num_local_experts
242 + w1_ep = w1[e_start:e_end]
243 + w2_ep = w2[e_start:e_end]
244 +
245 + w1_scale_ep, w2_scale_ep = None, None
246 + if is_quantized:
247 + w1_scale_ep = w1_scale[e_start:e_end] # type: ignore
248 + w2_scale_ep = w2_scale[e_start:e_end] # type: ignore
249 +
250 + deepep_combined = deepep_v2_moe_impl(
251 + pg,
252 + pgi,
253 + dp_size,
254 + test_tensors,
255 + w1_ep,
256 + w2_ep,
257 + w1_scale_ep,
258 + w2_scale_ep,
259 + config.num_experts,
260 + config.topk,
261 + use_fp8_dispatch,
262 + per_act_token_quant,
263 + )
264 +
265 + torch.testing.assert_close(
266 + torch_combined,
267 + deepep_combined,
268 + atol=6e-2,
269 + rtol=6e-2,
270 + )
271 +
272 +
273 +MNKs = [
274 + (1, 256, 256),
275 + (2, 256, 512),
276 + (3, 1024, 2048),
277 + (32, 256, 1024),
278 + (45, 512, 2048),
279 + (64, 1024, 1024),
280 + (222, 1024, 2048),
281 +]
282 +
283 +DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
284 +
285 +
286 +@pytest.mark.parametrize("dtype", DTYPES)
287 +@pytest.mark.parametrize("m,n,k", MNKs)
288 +@pytest.mark.parametrize("num_experts", [32])
289 +@pytest.mark.parametrize("topk", [6])
290 +@pytest.mark.parametrize("world_dp_size", [(2, 1)])
291 +@multi_gpu_test(num_gpus=2)
292 +@requires_deep_ep_v2
293 +def test_deep_ep_v2_moe(
294 + dtype: torch.dtype,
295 + m: int,
296 + n: int,
297 + k: int,
298 + num_experts: int,
299 + topk: int,
300 + world_dp_size: tuple[int, int],
301 + workspace_init,
302 +):
303 + per_act_token_quant = False
304 + use_fp8_dispatch = False
305 +
306 + set_random_seed(7)
307 + world_size, dp_size = world_dp_size
308 + config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts)
309 +
310 + quant_dtype = dtype if dtype == torch.float8_e4m3fn else None
311 + (_, w1, w1_scale, _), (_, w2, w2_scale, _) = make_test_weights(
312 + num_experts,
313 + n,
314 + k,
315 + quant_dtype=quant_dtype,
316 + per_out_ch_quant=True,
317 + )
318 +
319 + parallel_launch(
320 + world_size,
321 + _deep_ep_v2_moe,
322 + dp_size,
323 + config,
324 + w1,
325 + w2,
326 + w1_scale,
327 + w2_scale,
328 + use_fp8_dispatch,
329 + per_act_token_quant,
330 + )
331 +
332 +
333 +def _deep_ep_v2_moe_cudagraph(
334 + pgi: ProcessGroupInfo,
335 + dp_size: int,
336 + config: TestConfig,
337 + w1: torch.Tensor,
338 + w2: torch.Tensor,
339 + w1_scale: torch.Tensor | None,
340 + w2_scale: torch.Tensor | None,
341 +):
342 + """Worker function: verify DeepEP v2 + TrtLLM FP8 with do_expand=False."""
343 + import tempfile
344 +
345 + from vllm.distributed import (
346 + init_distributed_environment,
347 + initialize_model_parallel,
348 + )
349 +
350 + device = torch.device(f"cuda:{pgi.local_rank}")
351 + init_workspace_manager(device)
352 +
353 + pg = torch.distributed.new_group(list(range(pgi.world_size)))
354 + test_tensors = TestTensors.make(config)
355 + num_local_experts = config.num_experts // pgi.world_size
356 + hidden_size = config.k
357 +
358 + # Create FP8 weights directly, then dequantize for bf16 reference.
359 + w1_fp8 = torch.randn(
360 + (config.num_experts, 2 * config.n, config.k),
361 + device="cuda",
362 + dtype=torch.bfloat16,
363 + ).to(torch.float8_e4m3fn)
364 + w2_fp8 = torch.randn(
365 + (config.num_experts, config.k, config.n),
366 + device="cuda",
367 + dtype=torch.bfloat16,
368 + ).to(torch.float8_e4m3fn)
369 + w1_ref = w1_fp8.to(torch.bfloat16)
370 + w2_ref = w2_fp8.to(torch.bfloat16)
371 +
372 + from vllm.config import KernelConfig
373 +
374 + vllm_cfg = VllmConfig()
375 + vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm")
376 +
377 + with set_current_vllm_config(vllm_cfg):
378 + # Initialize vLLM parallel state (needed by FusedMoE layer)
379 + temp_file = tempfile.mktemp()
380 + init_distributed_environment(
381 + world_size=pgi.world_size,
382 + rank=pgi.rank,
383 + distributed_init_method=f"file://{temp_file}",
384 + local_rank=pgi.local_rank,
385 + backend="nccl",
386 + )
387 + initialize_model_parallel(tensor_model_parallel_size=1)
388 + # Reference MoE using dequantized bf16 weights
389 + torch_combined = torch_experts(
390 + test_tensors.rank_tokens,
391 + w1_ref,
392 + w2_ref,
393 + test_tensors.topk_weights,
394 + test_tensors.topk,
395 + )
396 +
397 + # Use the production pipeline: make_fused_moe_layer creates
398 + # a FusedMoE layer, quantizes weights, runs
399 + # process_weights_after_loading (TrtLLM W31 swap + BlockMajorK
400 + # shuffle), and selects the kernel.
401 + # Quantize weights using production helper, EP-slice, then
402 + # convert to TrtLLM format.
403 + from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves
404 + from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
405 + TrtLlmFp8ExpertsModular,
406 + )
407 + from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
408 + Fp8MoeBackend,
409 + convert_to_fp8_moe_kernel_format,
410 + )
411 +
412 + block_shape = [128, 128]
413 + qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape)
414 +
415 + # EP-slice before format conversion
416 + e_start = num_local_experts * pgi.rank
417 + e_end = e_start + num_local_experts
418 + w1_ep = qw.w13_weight[e_start:e_end]
419 + w2_ep = qw.w2_weight[e_start:e_end]
420 + assert qw.w13_weight_scale is not None
421 + assert qw.w2_weight_scale is not None
422 + w1_scale_ep = qw.w13_weight_scale[e_start:e_end]
423 + w2_scale_ep = qw.w2_weight_scale[e_start:e_end]
424 +
425 + # Convert to TrtLLM format (W31 swap + BlockMajorK shuffle)
426 + class _MockLayer:
427 + weight_block_size = block_shape
428 +
429 + class moe_config:
430 + is_act_and_mul = True
431 + intermediate_size_per_partition = config.n
432 +
433 + class activation:
434 + is_gated = True
435 +
436 + w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = convert_to_fp8_moe_kernel_format(
437 + fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM,
438 + layer=_MockLayer(),
439 + w13=w1_ep,
440 + w2=w2_ep,
441 + w13_scale=w1_scale_ep,
442 + w2_scale=w2_scale_ep,
443 + w13_input_scale=None,
444 + w2_input_scale=None,
445 + )
446 +
447 + # Build TrtLLM expert with correct EP params
448 + quant_config = FusedMoEQuantConfig.make(
449 + torch.float8_e4m3fn,
450 + block_shape=block_shape,
451 + w1_scale=w1_scale_ep,
452 + w2_scale=w2_scale_ep,
453 + )
454 + moe_config = make_dummy_moe_config(
455 + num_experts=num_local_experts,
456 + experts_per_token=config.topk,
457 + hidden_dim=hidden_size,
458 + intermediate_size=config.n,
459 + )
460 + fused_experts = TrtLlmFp8ExpertsModular(
461 + moe_config=moe_config,
462 + quant_config=quant_config,
463 + )
464 +
465 + v2_args = DeepEPV2Args(
466 + num_local_experts=num_local_experts,
467 + num_experts=config.num_experts,
468 + num_topk=config.topk,
469 + hidden_size=hidden_size,
470 + max_tokens_per_rank=8192,
471 + use_fp8_dispatch=False,
472 + )
473 + a2a = make_deepep_v2_a2a(
474 + pg=pg,
475 + pgi=pgi,
476 + dp_size=dp_size,
477 + v2_args=v2_args,
478 + use_cudagraph=True,
479 + )
480 + mk_kernel = FusedMoEKernel(
481 + prepare_finalize=a2a,
482 + fused_experts=fused_experts,
483 + inplace=False,
484 + )
485 +
486 + for _ in range(3):
487 + out = mk_kernel.apply(
488 + hidden_states=test_tensors.rank_tokens,
489 + w1=w1_ep,
490 + w2=w2_ep,
491 + topk_weights=test_tensors.topk_weights,
492 + topk_ids=test_tensors.topk,
493 + activation=MoEActivation.SILU,
494 + global_num_experts=config.num_experts,
495 + expert_map=None,
496 + apply_router_weight_on_input=False,
497 + )
498 +
499 + torch.testing.assert_close(
500 + torch_combined,
501 + out,
502 + atol=6e-2,
503 + rtol=6e-2,
504 + )
505 +
506 +
507 +@pytest.mark.parametrize("m,n,k", [(32, 256, 1024)])
508 +@pytest.mark.parametrize("num_experts", [32])
509 +@pytest.mark.parametrize("topk", [6])
510 +@pytest.mark.parametrize("world_dp_size", [(2, 1)])
511 +@multi_gpu_test(num_gpus=2)
512 +@requires_deep_ep_v2
513 +def test_deep_ep_v2_moe_cudagraph(
514 + m: int,
515 + n: int,
516 + k: int,
517 + num_experts: int,
518 + topk: int,
519 + world_dp_size: tuple[int, int],
520 + workspace_init,
521 +):
522 + set_random_seed(7)
523 + world_size, dp_size = world_dp_size
524 + config = TestConfig(
525 + dtype=torch.float8_e4m3fn,
526 + topk=topk,
527 + m=m,
528 + k=k,
529 + n=n,
530 + num_experts=num_experts,
531 + )
532 +
533 + parallel_launch(
534 + world_size,
535 + _deep_ep_v2_moe_cudagraph,
536 + dp_size,
537 + config,
538 + None, # weights created inside worker
539 + None,
540 + None,
541 + None,
542 + )
tools/ep_kernels/install_python_libraries.sh
+2 -1
@@ -8,7 +8,8 @@ set -ex
8 # --nvshmem-ver <ver> NVSHMEM version
9
10 CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
11 -DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
11 +DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"d4f41e4e93"}
12 +
13 NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"} # Default supports both CUDA 12 and 13
14 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
15 MODE=${MODE:-install}
vllm/config/parallel.py
+1
@@ -42,6 +42,7 @@ All2AllBackend = Literal[
42 "pplx",
43 "deepep_high_throughput",
44 "deepep_low_latency",
45 + "deepep_v2",
46 "mori_high_throughput",
47 "mori_low_latency",
48 "nixl_ep",
vllm/distributed/device_communicators/all2all.py
+64 -1
@@ -15,7 +15,7 @@ from vllm.utils.flashinfer import (
15 has_flashinfer_nvlink_one_sided,
16 has_flashinfer_nvlink_two_sided,
17 )
18 -from vllm.utils.import_utils import has_deep_ep, has_mori
18 +from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori
19
20 from .base_device_communicator import All2AllManagerBase, Cache
21
@@ -863,3 +863,66 @@ class MoriAll2AllManager(All2AllManagerBase):
863 mori_kwargs, self._make_handle
864 )
865 return handle
866 +
867 +
868 +class DeepEPV2All2AllManager(All2AllManagerBase):
869 + """
870 + All2All communication based on DeepEP v2 ElasticBuffer (unified API).
871 + Uses NCCL Gin backend with analytical SM calculation.
872 + """
873 +
874 + def __init__(self, cpu_group, tcp_store_group=None, device_group=None):
875 + assert has_deep_ep_v2(), (
876 + "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 "
877 + "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4."
878 + )
879 + super().__init__(cpu_group, tcp_store_group)
880 + self._device_group = device_group
881 + self.handle_cache = Cache()
882 + self._num_sms: int | None = None
883 +
884 + def _make_all2all_kwargs(
885 + self,
886 + num_max_tokens_per_rank: int,
887 + hidden: int,
888 + num_topk: int,
889 + use_fp8_dispatch: bool,
890 + ) -> dict:
891 + return dict(
892 + group=self._device_group
893 + if self._device_group is not None
894 + else self.cpu_group,
895 + num_max_tokens_per_rank=num_max_tokens_per_rank,
896 + hidden=hidden,
897 + num_topk=num_topk,
898 + use_fp8_dispatch=use_fp8_dispatch,
899 + allow_hybrid_mode=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE,
900 + prefer_overlap_with_compute=envs.VLLM_DEEPEP_V2_PREFER_OVERLAP,
901 + allow_multiple_reduction=(envs.VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION),
902 + explicitly_destroy=True,
903 + )
904 +
905 + def get_handle(self, kwargs):
906 + import deep_ep # type: ignore[import-not-found]
907 +
908 + num_experts = kwargs.pop("num_experts", 256)
909 + buffer_kwargs = self._make_all2all_kwargs(**kwargs)
910 + logger.debug("DeepEP v2 all2all args %s", buffer_kwargs)
911 + handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create(
912 + buffer_kwargs, deep_ep.ElasticBuffer
913 + )
914 + if self._num_sms is None:
915 + self._num_sms = handle.get_theoretical_num_sms(
916 + num_experts=num_experts,
917 + num_topk=kwargs["num_topk"],
918 + )
919 + return handle
920 +
921 + def max_sms_used(self) -> int | None:
922 + return self._num_sms
923 +
924 + def destroy(self):
925 + with self.handle_cache._lock:
926 + for _, handle in self.handle_cache._cache.items():
927 + handle.destroy()
928 + self.handle_cache._cache.clear()
vllm/distributed/device_communicators/cuda_communicator.py
+8
@@ -146,6 +146,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
146 self.all2all_manager = MoriAll2AllManager(
147 self.cpu_group, self.all2all_backend
148 )
149 + elif self.all2all_backend == "deepep_v2":
150 + from .all2all import DeepEPV2All2AllManager
151 +
152 + self.all2all_manager = DeepEPV2All2AllManager(
153 + self.cpu_group,
154 + tcp_store_group,
155 + device_group=self.device_group,
156 + )
157 elif self.all2all_backend == "nixl_ep":
158 from .all2all import NixlEPAll2AllManager
159
vllm/envs.py
+15
@@ -246,6 +246,9 @@ if TYPE_CHECKING:
246 VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024
247 VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: bool = False
248 VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL: bool = False
249 + VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE: bool = True
250 + VLLM_DEEPEP_V2_PREFER_OVERLAP: bool = False
251 + VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION: bool = False
252 VLLM_DBO_COMM_SMS: int = 20
253 VLLM_PATTERN_MATCH_DEBUG: str | None = None
254 VLLM_DEBUG_DUMP_PATH: str | None = None
@@ -1840,6 +1843,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
1843 "VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL": lambda: bool(
1844 int(os.getenv("VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL", "0"))
1845 ),
1846 + # DeepEP v2: enable two-tier NVLink+RDMA hybrid mode
1847 + "VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE": lambda: bool(
1848 + int(os.getenv("VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE", "0"))
1849 + ),
1850 + # DeepEP v2: use fewer SMs at slight throughput cost
1851 + "VLLM_DEEPEP_V2_PREFER_OVERLAP": lambda: bool(
1852 + int(os.getenv("VLLM_DEEPEP_V2_PREFER_OVERLAP", "0"))
1853 + ),
1854 + # DeepEP v2: trade precision for transfer size in combine
1855 + "VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION": lambda: bool(
1856 + int(os.getenv("VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION", "0"))
1857 + ),
1858 # The number of SMs/CUs to allocate for communication kernels when
1859 # running DBO; the rest will be allocated to compute.
1860 # Default: 20 on CUDA (SMs), 64 on ROCm (CUs).
vllm/model_executor/layers/fused_moe/all2all_utils.py
+43 -1
@@ -29,7 +29,12 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two
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
32 +from vllm.utils.import_utils import (
33 + has_deep_ep,
34 + has_deep_ep_v2,
35 + has_mori,
36 + has_nixl_ep,
37 +)
38
39 logger = init_logger(__name__)
40
@@ -40,6 +45,8 @@ if current_platform.is_cuda_alike():
45 DEEPEP_QUANT_BLOCK_SHAPE,
46 DeepEPLLPrepareAndFinalize,
47 )
48 + if has_deep_ep_v2():
49 + from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
50 if has_mori():
51 from .prepare_finalize.mori import MoriPrepareAndFinalize
52 if has_nixl_ep():
@@ -94,6 +101,11 @@ def maybe_roundup_layer_hidden_size(
101 hidden_size
102 )
103
104 + if moe_parallel_config.use_deepep_v2_kernels:
105 + hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size(
106 + hidden_size, act_dtype
107 + )
108 +
109 if moe_parallel_config.use_nixl_ep_kernels:
110 hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size(
111 hidden_size
@@ -193,6 +205,36 @@ def maybe_make_prepare_finalize(
205 physical_to_global=physical_to_global,
206 local_expert_global_ids=local_expert_global_ids,
207 )
208 + elif moe.use_deepep_v2_kernels:
209 + assert moe.dp_size == all2all_manager.dp_world_size
210 +
211 + use_fp8_dispatch = (
212 + quant_config is not None
213 + and quant_config.quant_dtype == current_platform.fp8_dtype()
214 + and quant_config.is_block_quantized
215 + )
216 + all_to_all_args = dict(
217 + num_max_tokens_per_rank=moe.max_num_tokens,
218 + hidden=moe.hidden_dim,
219 + num_topk=moe.experts_per_token,
220 + num_experts=moe.num_experts,
221 + use_fp8_dispatch=use_fp8_dispatch,
222 + )
223 + handle = all2all_manager.get_handle(all_to_all_args)
224 + vllm_config = get_current_vllm_config()
225 + use_cudagraph = not vllm_config.model_config.enforce_eager
226 +
227 + prepare_finalize = DeepEPV2PrepareAndFinalize(
228 + buffer=handle,
229 + num_dispatchers=all2all_manager.world_size,
230 + dp_size=all2all_manager.dp_world_size,
231 + rank_expert_offset=all2all_manager.rank * moe.num_local_experts,
232 + num_experts=moe.num_experts,
233 + num_topk=moe.experts_per_token,
234 + use_fp8_dispatch=use_fp8_dispatch,
235 + use_cudagraph=use_cudagraph,
236 + )
237 +
238 elif moe.use_mori_kernels:
239 assert quant_config is not None
240
vllm/model_executor/layers/fused_moe/config.py
+8
@@ -1070,6 +1070,10 @@ class FusedMoEParallelConfig:
1070 def use_nixl_ep_kernels(self):
1071 return self.use_all2all_kernels and self.all2all_backend == "nixl_ep"
1072
1073 + @property
1074 + def use_deepep_v2_kernels(self):
1075 + return self.use_all2all_kernels and self.all2all_backend == "deepep_v2"
1076 +
1077 @staticmethod
1078 def flatten_tp_across_dp_and_pcp(
1079 tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
@@ -1394,6 +1398,10 @@ class FusedMoEConfig:
1398 def use_nixl_ep_kernels(self):
1399 return self.moe_parallel_config.use_nixl_ep_kernels
1400
1401 + @property
1402 + def use_deepep_v2_kernels(self):
1403 + return self.moe_parallel_config.use_deepep_v2_kernels
1404 +
1405 @property
1406 def needs_round_robin_routing_tables(self):
1407 return self.moe_parallel_config.needs_round_robin_routing_tables
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+9 -1
@@ -94,6 +94,14 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
94 Fp8 TRTLLM-Gen MoE kernels. Supports modular interface.
95 """
96
97 + @staticmethod
98 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
99 + return (
100 + not moe_parallel_config.use_all2all_kernels
101 + or moe_parallel_config.use_ag_rs_all2all_kernels
102 + or moe_parallel_config.use_deepep_v2_kernels
103 + ) and not moe_parallel_config.enable_eplb
104 +
105 @staticmethod
106 def _supports_quant_scheme(
107 weight_key: QuantKey | None,
@@ -193,7 +201,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
201 gemm2_weights=w2,
202 gemm2_weights_scale=self.quant_config.w2_scale,
203 num_experts=global_num_experts,
196 - top_k=self.topk,
204 + top_k=topk_ids.size(1),
205 n_group=None,
206 topk_group=None,
207 intermediate_size=self.intermediate_size_per_partition,
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+12
@@ -85,6 +85,18 @@ def _get_priority_backends(
85 def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None:
86 backends.insert(0, backends.pop(backends.index(backend)))
87
88 + # With DeepEP v2 contiguous layout (do_expand=False), tensors are
89 + # worst-case allocated with padding. TrtLLM's tile-level skipping
90 + # avoids wasted compute on padding rows; other backends process all rows.
91 + if (
92 + current_platform.is_cuda()
93 + and current_platform.is_device_capability_family(100)
94 + and moe_config.moe_parallel_config.use_deepep_v2_kernels
95 + and activation_key == kFp8Dynamic128Sym
96 + and weight_key == kFp8Static128BlockSym
97 + ):
98 + _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.FLASHINFER_TRTLLM)
99 +
100 # On Hopper for Block Fp8, prefer Triton for TP and FI CUTLASS for EP.
101 if (
102 current_platform.is_cuda()
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py new
+394
@@ -0,0 +1,394 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 +
5 +import deep_ep
6 +import torch
7 +
8 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
11 + TopKWeightAndReduceContiguous,
12 + TopKWeightAndReduceDelegate,
13 +)
14 +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
15 +from vllm.utils.math_utils import round_up
16 +from vllm.v1.worker.ubatching import (
17 + dbo_current_ubatch_id,
18 +)
19 +
20 +
21 +class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
22 + """
23 + Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API).
24 +
25 + Supports two modes controlled by the `use_cudagraph` constructor arg:
26 +
27 + **Decode mode (use_cudagraph=True):**
28 + - do_expand=False, do_cpu_sync=False
29 + - Tokens returned in original order with recv_topk_idx (global IDs)
30 + - Worst-case tensor allocation; padding rows zeroed via
31 + handle.psum_num_recv_tokens_per_scaleup_rank
32 + - Fully cudagraph-capturable
33 + - Expert kernel sorts internally (expert_tokens_meta=None)
34 +
35 + **Prefill mode (use_cudagraph=False):**
36 + - do_expand=True, do_cpu_sync=True
37 + - Per-expert-contiguous layout; exact memory allocation
38 + - Saves GPU memory (no worst-case allocation)
39 + - Not cudagraph-capturable (CPU polling), but prefill doesn't
40 + use cudagraphs anyway
41 + - Provides expert_tokens_meta for efficient batched expert kernels
42 +
43 + Both modes use async_with_compute_stream=False (synchronous from
44 + caller's perspective). The ElasticBuffer handles comm internally.
45 + """
46 +
47 + @staticmethod
48 + def maybe_roundup_layer_hidden_size(hidden_size: int, dtype: torch.dtype) -> int:
49 + hidden_size_bytes = hidden_size * dtype.itemsize
50 + xfer_atom_size = 512 # 32 * 16 (size(int4))
51 + if hidden_size_bytes % xfer_atom_size == 0:
52 + return hidden_size
53 +
54 + hidden_size_bytes = round_up(hidden_size_bytes, xfer_atom_size)
55 + return hidden_size_bytes // dtype.itemsize
56 +
57 + def __init__(
58 + self,
59 + buffer: deep_ep.ElasticBuffer,
60 + num_dispatchers: int,
61 + dp_size: int,
62 + rank_expert_offset: int,
63 + num_experts: int,
64 + num_topk: int,
65 + use_fp8_dispatch: bool = False,
66 + use_cudagraph: bool = False,
67 + ):
68 + super().__init__()
69 + self.buffer = buffer
70 + self.num_dispatchers_ = num_dispatchers
71 + self.dp_size = dp_size
72 + self.rank_expert_offset = rank_expert_offset
73 + self.num_experts = num_experts
74 + self.num_topk = num_topk
75 + self.use_fp8_dispatch = use_fp8_dispatch
76 + self.use_cudagraph = use_cudagraph
77 +
78 + # DBO microbatching: one handle slot per micro-batch.
79 + self.handles: list[deep_ep.EPHandle | None] = [None, None]
80 +
81 + def num_dispatchers(self) -> int:
82 + return self.num_dispatchers_
83 +
84 + def output_is_reduced(self) -> bool:
85 + return True
86 +
87 + @property
88 + def activation_format(self) -> mk.FusedMoEActivationFormat:
89 + return mk.FusedMoEActivationFormat.Standard
90 +
91 + def max_num_tokens_per_rank(self) -> int | None:
92 + return None
93 +
94 + def topk_indices_dtype(self) -> torch.dtype | None:
95 + return torch.int64
96 +
97 + def _do_dispatch(
98 + self,
99 + tokens: torch.Tensor,
100 + token_scales: torch.Tensor | None,
101 + rank_topk_ids: torch.Tensor,
102 + rank_topk_weights: torch.Tensor,
103 + num_experts: int,
104 + a1_scale: torch.Tensor | None,
105 + quant_config: FusedMoEQuantConfig,
106 + defer_input_quant: bool,
107 + ) -> Callable:
108 + has_scales = token_scales is not None
109 +
110 + token_data = tokens
111 + if has_scales:
112 + token_data = (tokens, token_scales)
113 +
114 + # Decode: do_expand=False + do_cpu_sync=False (cudagraph-safe)
115 + # Prefill: do_expand=True + do_cpu_sync=True (memory-efficient)
116 + do_expand = not self.use_cudagraph
117 + do_cpu_sync = not self.use_cudagraph
118 +
119 + (
120 + recv_x,
121 + recv_topk_idx,
122 + recv_topk_weights,
123 + handle,
124 + event,
125 + ) = self.buffer.dispatch(
126 + x=token_data,
127 + topk_idx=rank_topk_ids,
128 + topk_weights=rank_topk_weights,
129 + num_experts=num_experts,
130 + do_expand=do_expand,
131 + do_cpu_sync=do_cpu_sync,
132 + async_with_compute_stream=False,
133 + )
134 +
135 + a2a_idx = dbo_current_ubatch_id()
136 + self.handles[a2a_idx] = handle
137 +
138 + return lambda: self._receiver(
139 + event,
140 + has_scales,
141 + recv_x,
142 + recv_topk_idx,
143 + num_experts,
144 + handle.num_recv_tokens_per_expert_list,
145 + recv_topk_weights,
146 + handle.psum_num_recv_tokens_per_scaleup_rank,
147 + a1_scale,
148 + quant_config,
149 + defer_input_quant=defer_input_quant,
150 + )
151 +
152 + def _receiver(
153 + self,
154 + event: deep_ep.EventOverlap,
155 + has_scales: bool,
156 + recv_x: tuple[torch.Tensor, torch.Tensor] | torch.Tensor,
157 + recv_topk_idx: torch.Tensor | None,
158 + num_experts: int,
159 + recv_expert_num_tokens: list[int],
160 + recv_topk_weights: torch.Tensor | None,
161 + psum_recv_per_rank: torch.Tensor,
162 + a1_scale: torch.Tensor | None,
163 + quant_config: FusedMoEQuantConfig,
164 + defer_input_quant: bool,
165 + ) -> mk.PrepareResultType:
166 + if event.event is not None:
167 + event.current_stream_wait()
168 +
169 + if isinstance(recv_x, tuple):
170 + expert_x, expert_x_scale = recv_x
171 + else:
172 + expert_x, expert_x_scale = recv_x, None
173 +
174 + if recv_topk_idx is None:
175 + # do_expand=True (prefill mode): build topk_ids from
176 + # per-expert token counts.
177 + total_tokens = sum(recv_expert_num_tokens)
178 + if total_tokens > 0:
179 + recv_topk_idx = torch.empty(
180 + total_tokens,
181 + dtype=torch.int64,
182 + device=expert_x.device,
183 + )
184 + offset = 0
185 + for i, count in enumerate(recv_expert_num_tokens):
186 + if count > 0:
187 + recv_topk_idx[offset : offset + count].fill_(
188 + i + self.rank_expert_offset
189 + )
190 + offset += count
191 + else:
192 + recv_topk_idx = torch.empty(
193 + 0,
194 + dtype=torch.int64,
195 + device=expert_x.device,
196 + )
197 + recv_topk_idx = recv_topk_idx.unsqueeze(1)
198 + else:
199 + # do_expand=False (decode/cudagraph mode): recv_topk_idx has
200 + # LOCAL expert IDs (-1 for non-local and padding rows).
201 + # Convert valid local IDs to global. Rows with -1 are
202 + # skipped by expert kernels (TrtLLM tile-level skipping,
203 + # DeepGemm is_computation_valid), so no need to zero
204 + # hidden states, scales, or weights for padding rows.
205 + valid_mask = recv_topk_idx >= 0
206 + recv_topk_idx = torch.where(
207 + valid_mask,
208 + recv_topk_idx + self.rank_expert_offset,
209 + recv_topk_idx,
210 + )
211 +
212 + # Reshape recv_topk_weights to match recv_topk_idx shape [N, 1]
213 + if recv_topk_weights is not None and recv_topk_weights.ndim == 1:
214 + recv_topk_weights = recv_topk_weights.unsqueeze(1)
215 +
216 + expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
217 + recv_expert_num_tokens,
218 + device=expert_x.device,
219 + )
220 +
221 + if not quant_config.is_block_quantized and not defer_input_quant:
222 + expert_x_scale = None
223 + if expert_x.numel() != 0:
224 + expert_x, expert_x_scale = moe_kernel_quantize_input(
225 + expert_x,
226 + a1_scale,
227 + quant_dtype=quant_config.quant_dtype,
228 + per_act_token_quant=False,
229 + block_shape=quant_config.block_shape,
230 + is_scale_swizzled=quant_config.is_scale_swizzled,
231 + )
232 +
233 + return (
234 + expert_x,
235 + expert_x_scale,
236 + expert_tokens_meta,
237 + recv_topk_idx,
238 + recv_topk_weights,
239 + )
240 +
241 + def supports_async(self) -> bool:
242 + return True
243 +
244 + def prepare_async(
245 + self,
246 + a1: torch.Tensor,
247 + topk_weights: torch.Tensor,
248 + topk_ids: torch.Tensor,
249 + num_experts: int,
250 + expert_map: torch.Tensor | None,
251 + apply_router_weight_on_input: bool,
252 + quant_config: FusedMoEQuantConfig,
253 + defer_input_quant: bool = False,
254 + ) -> mk.ReceiverType:
255 + if apply_router_weight_on_input:
256 + topk = topk_ids.size(1)
257 + assert topk == 1, (
258 + "apply_router_weight_on_input is only implemented for topk=1"
259 + )
260 + a1 = a1 * topk_weights.to(a1.dtype)
261 +
262 + if quant_config.is_block_quantized and not defer_input_quant:
263 + a1q, a1q_scale = moe_kernel_quantize_input(
264 + a1,
265 + quant_config.a1_scale,
266 + quant_dtype=quant_config.quant_dtype,
267 + per_act_token_quant=quant_config.per_act_token_quant,
268 + block_shape=quant_config.block_shape,
269 + )
270 + if a1q_scale is not None and a1q_scale.numel() == 1:
271 + a1q_scale = a1q_scale.view(1, 1)
272 + a1_post_scale = None
273 + else:
274 + a1q = a1
275 + a1q_scale = None
276 + a1_post_scale = (
277 + quant_config.a1_gscale
278 + if quant_config.quant_dtype == "nvfp4"
279 + else quant_config.a1_scale
280 + )
281 +
282 + return self._do_dispatch(
283 + tokens=a1q,
284 + token_scales=a1q_scale,
285 + rank_topk_ids=topk_ids,
286 + rank_topk_weights=topk_weights,
287 + num_experts=num_experts,
288 + a1_scale=a1_post_scale,
289 + quant_config=quant_config,
290 + defer_input_quant=defer_input_quant,
291 + )
292 +
293 + def prepare(
294 + self,
295 + a1: torch.Tensor,
296 + topk_weights: torch.Tensor,
297 + topk_ids: torch.Tensor,
298 + num_experts: int,
299 + expert_map: torch.Tensor | None,
300 + apply_router_weight_on_input: bool,
301 + quant_config: FusedMoEQuantConfig,
302 + defer_input_quant: bool = False,
303 + ) -> mk.PrepareResultType:
304 + receiver = self.prepare_async(
305 + a1,
306 + topk_weights,
307 + topk_ids,
308 + num_experts,
309 + expert_map,
310 + apply_router_weight_on_input,
311 + quant_config,
312 + defer_input_quant,
313 + )
314 + return receiver()
315 +
316 + def _finalize(
317 + self,
318 + output: torch.Tensor,
319 + fused_expert_output: torch.Tensor,
320 + topk_weights: torch.Tensor,
321 + topk_ids: torch.Tensor,
322 + apply_router_weight_on_input: bool,
323 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
324 + do_async: bool,
325 + ) -> Callable | None:
326 + a2a_idx = dbo_current_ubatch_id()
327 + handle = self.handles[a2a_idx]
328 + assert handle is not None
329 +
330 + if fused_expert_output.numel() != 0:
331 + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
332 + weight_and_reduce_impl = TopKWeightAndReduceContiguous()
333 + fused_expert_output = weight_and_reduce_impl.apply(
334 + output=None,
335 + fused_expert_output=fused_expert_output,
336 + topk_weights=topk_weights,
337 + topk_ids=topk_ids,
338 + apply_router_weight_on_input=apply_router_weight_on_input,
339 + )
340 +
341 + if fused_expert_output.dtype != torch.bfloat16:
342 + raise ValueError(
343 + f"DeepEP v2 combine requires bfloat16 input, "
344 + f"got {fused_expert_output.dtype}"
345 + )
346 +
347 + combined_x, _, event = self.buffer.combine(
348 + x=fused_expert_output,
349 + handle=handle,
350 + topk_weights=None,
351 + async_with_compute_stream=False,
352 + )
353 +
354 + output.copy_(combined_x, non_blocking=True)
355 + return None
356 +
357 + def finalize_async(
358 + self,
359 + output: torch.Tensor,
360 + fused_expert_output: torch.Tensor,
361 + topk_weights: torch.Tensor,
362 + topk_ids: torch.Tensor,
363 + apply_router_weight_on_input: bool,
364 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
365 + ) -> Callable:
366 + self._finalize(
367 + output,
368 + fused_expert_output,
369 + topk_weights,
370 + topk_ids,
371 + apply_router_weight_on_input,
372 + weight_and_reduce_impl,
373 + False,
374 + )
375 + return lambda: None
376 +
377 + def finalize(
378 + self,
379 + output: torch.Tensor,
380 + fused_expert_output: torch.Tensor,
381 + topk_weights: torch.Tensor,
382 + topk_ids: torch.Tensor,
383 + apply_router_weight_on_input: bool,
384 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
385 + ) -> None:
386 + self._finalize(
387 + output,
388 + fused_expert_output,
389 + topk_weights,
390 + topk_ids,
391 + apply_router_weight_on_input,
392 + weight_and_reduce_impl,
393 + False,
394 + )
vllm/model_executor/layers/fused_moe/routed_experts.py
+1
@@ -1038,6 +1038,7 @@ class RoutedExperts(PluggableLayer):
1038 "e_score_correction_bias",
1039 "w13_input_scale",
1040 "w2_input_scale",
1041 + "hash_indices_table",
1042 }
1043
1044 # Parameters of non-expert submodules that live inside runner (RoutedExperts).
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+8
@@ -168,6 +168,14 @@ def fused_topk_bias(
168 hash_indices_table: torch.Tensor | None = None,
169 routed_scaling_factor: float = 1.0,
170 ):
171 + # The topk kernel dispatches dtype based on topk_ids (set by
172 + # indices_type) and assumes input_tokens/hash_indices_table match.
173 + if indices_type is not None:
174 + if input_tokens is not None and input_tokens.dtype != indices_type:
175 + input_tokens = input_tokens.to(dtype=indices_type)
176 + if hash_indices_table is not None and hash_indices_table.dtype != indices_type:
177 + hash_indices_table = hash_indices_table.to(dtype=indices_type)
178 +
179 if not rocm_aiter_ops.is_fused_moe_enabled():
180 assert hidden_states.size(0) == gating_output.size(0), (
181 "Number of tokens mismatch"
vllm/model_executor/layers/quantization/mxfp4.py
+4
@@ -489,6 +489,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
489 self.w13_precision_config = None
490 self.w2_precision_config = None
491
492 + @property
493 + def supports_eplb(self) -> bool:
494 + return True
495 +
496 @property
497 def skip_forward_padding(self) -> bool:
498 # SM100_FI_MXFP4_MXFP8_TRTLLM supports padding with mxfp8 quant
vllm/models/deepseek_v4/nvidia/model.py
+16 -9
@@ -566,7 +566,7 @@ class DeepseekV4MoE(nn.Module):
566 if self.use_mega_moe:
567 self._init_mega_moe_experts(vllm_config, config, prefix)
568 else:
569 - self._init_fused_moe_experts(config, quant_config, prefix)
569 + self._init_fused_moe_experts(vllm_config, config, quant_config, prefix)
570
571 def _init_mega_moe_experts(
572 self,
@@ -612,22 +612,27 @@ class DeepseekV4MoE(nn.Module):
612
613 def _init_fused_moe_experts(
614 self,
615 + vllm_config: VllmConfig,
616 config,
617 quant_config,
618 prefix: str,
619 ) -> None:
620 + parallel_config = vllm_config.parallel_config
621 self.tp_rank = get_tensor_model_parallel_rank()
620 - assert config.n_routed_experts % self.tp_size == 0
621 -
622 - self.n_local_experts = config.n_routed_experts // self.tp_size
623 - self.experts_start_idx = self.tp_rank * self.n_local_experts
624 - self.experts_end_idx = self.experts_start_idx + self.n_local_experts
622
626 - self.n_redundant_experts = 0
623 + eplb_config = parallel_config.eplb_config
624 + self.n_redundant_experts = eplb_config.num_redundant_experts
625 self.n_shared_experts = config.n_shared_experts or 0
626 self.n_logical_experts = self.n_routed_experts
629 - self.n_physical_experts = self.n_logical_experts
630 - self.n_local_physical_experts = self.n_local_experts
627 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
628 + assert self.n_physical_experts % self.tp_size == 0, (
629 + f"n_physical_experts={self.n_physical_experts} must be divisible by "
630 + f"tp_size={self.tp_size}. Adjust num_redundant_experts."
631 + )
632 + self.n_local_physical_experts = self.n_physical_experts // self.tp_size
633 + self.n_local_experts = self.n_local_physical_experts
634 + self.experts_start_idx = self.tp_rank * self.n_local_experts
635 + self.experts_end_idx = self.experts_start_idx + self.n_local_experts
636 self.physical_expert_start = self.experts_start_idx
637 self.physical_expert_end = self.experts_end_idx
638
@@ -647,6 +652,8 @@ class DeepseekV4MoE(nn.Module):
652 hash_indices_table=self.gate.tid2eid,
653 swiglu_limit=self.swiglu_limit,
654 router_logits_dtype=torch.float32,
655 + enable_eplb=parallel_config.enable_eplb,
656 + num_redundant_experts=eplb_config.num_redundant_experts,
657 )
658
659 def forward(
vllm/utils/import_utils.py
+55
@@ -417,6 +417,61 @@ def has_deep_ep() -> bool:
417 return _has_module("deep_ep")
418
419
420 +DEEPEP_V2_MIN_NCCL_VERSION_RAW = 23004 # 2.30.4
421 +
422 +
423 +def _get_runtime_nccl_version() -> int | None:
424 + """Get the runtime NCCL version by loading the actual library.
425 +
426 + Returns the raw version int (e.g. 23004 for 2.30.4), or None on failure.
427 + torch.cuda.nccl.version() is a compile-time constant from the PyTorch
428 + wheel and does not reflect a separately installed NCCL.
429 + """
430 + import ctypes
431 +
432 + try:
433 + from vllm.utils.nccl import find_nccl_library
434 +
435 + lib = ctypes.CDLL(find_nccl_library())
436 + version = ctypes.c_int()
437 + lib.ncclGetVersion(ctypes.byref(version))
438 + return version.value
439 + except Exception:
440 + return None
441 +
442 +
443 +def _format_nccl_raw_version(raw: int) -> str:
444 + s = str(raw)
445 + return f"{s[0]}.{s[1:3].lstrip('0') or '0'}.{s[3:].lstrip('0') or '0'}"
446 +
447 +
448 +def has_deep_ep_v2() -> bool:
449 + """Whether deep_ep with ElasticBuffer (v2 API) is available.
450 +
451 + Requires both the ElasticBuffer class in the deep_ep module and
452 + NCCL >= 2.30.4 (GIN backend), checked against the runtime library.
453 + """
454 + if not _has_module("deep_ep"):
455 + return False
456 + import deep_ep # type: ignore[import-not-found]
457 +
458 + if not hasattr(deep_ep, "ElasticBuffer"):
459 + return False
460 + try:
461 + nccl_ver = _get_runtime_nccl_version()
462 + if nccl_ver is None or nccl_ver < DEEPEP_V2_MIN_NCCL_VERSION_RAW:
463 + logger.info_once(
464 + "DeepEP v2 requires NCCL >= %s but found %s. "
465 + "deepep_v2 backend will not be available.",
466 + _format_nccl_raw_version(DEEPEP_V2_MIN_NCCL_VERSION_RAW),
467 + _format_nccl_raw_version(nccl_ver) if nccl_ver else "unknown",
468 + )
469 + return False
470 + except Exception:
471 + return False
472 + return True
473 +
474 +
475 def has_deep_gemm() -> bool:
476 """Whether the optional `deep_gemm` package is available.
477