@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 911548247a

[Distributed][MoonEP] BF16 integration of MoonEP balanced EP backend (#52101)

Signed-off-by: kevinli <306243285+kaijunli-infr@users.noreply.github.com> Co-authored-by: kevinli <306243285+kaijunli-infr@users.noreply.github.com>

Kevin Li committed Sep 18, 2026 at 07:34 UTC 911548247a7b04e0c9cffca1cfba66b7120d6d3f
15 files changed +1338 -12
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-BF16-moonep.yaml new
+5
@@ -0,0 +1,5 @@
1 +model_name: "Qwen/Qwen3-30B-A3B"
2 +accuracy_threshold: 0.88
3 +num_questions: 1319
4 +num_fewshot: 5
5 +server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend=moonep"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
+1
@@ -10,3 +10,4 @@ Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ht.yaml
10 Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml
11 Qwen3-30B-A3B-Fp8-CT-Block-deepgemm.yaml
12 Qwen3-30B-A3B-BF16-triton.yaml
13 +Qwen3-30B-A3B-BF16-moonep.yaml
tests/evals/gsm8k/test_gsm8k_correctness.py
+6
@@ -128,6 +128,12 @@ def test_gsm8k_correctness(config_filename):
128 "Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. "
129 "The quantization scheme is not supported on non-GFX950 platforms."
130 )
131 + if "--all2all-backend=moonep" in eval_config.get("server_args", ""):
132 + from vllm.utils.import_utils import has_moonep
133 +
134 + if not has_moonep():
135 + pytest.skip("Skipping MoonEP config: the moonep package is not installed.")
136 +
137 # Parse server arguments from config (use shlex to handle quoted strings)
138 server_args_str = eval_config.get("server_args", "")
139 server_args = shlex.split(server_args_str) if server_args_str else []
tests/kernels/moe/test_moonep_bf16_poc.py new
+412
@@ -0,0 +1,412 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Test MoonEP dispatch / prefetch / combine logic (BF16 PoC).
4 +
5 +Two paths, both compared against the pure-PyTorch reference MoE:
6 +- MoonEPPrepareAndFinalize + a reference segment-loop expert runner over
7 + MoonEP's expert-grouped ``[NvS, H]`` layout;
8 +- the full modular kernel, MoonEPPrepareAndFinalize + MoonEPExperts through
9 + FusedMoEKernel.
10 +Requires NVSwitch multicast capable GPUs.
11 +"""
12 +
13 +import dataclasses
14 +import types
15 +
16 +import pytest
17 +import torch
18 +import torch.nn.functional as F
19 +from torch.distributed import ProcessGroup
20 +
21 +from tests.kernels.moe.utils import make_test_weights
22 +from tests.kernels.utils import torch_experts
23 +from vllm.config import VllmConfig, set_current_vllm_config
24 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
25 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
26 + TopKWeightAndReduceNoOP,
27 +)
28 +from vllm.utils.import_utils import has_moonep
29 +from vllm.utils.torch_utils import set_random_seed
30 +
31 +from ...utils import multi_gpu_test
32 +from .parallel_utils import ProcessGroupInfo, parallel_launch
33 +
34 +if has_moonep():
35 + from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
36 + MoonEPExpertWeightLayout,
37 + MoonEPPrepareAndFinalize,
38 + make_moonep_weight_layout,
39 + )
40 +
41 +requires_moonep = pytest.mark.skipif(
42 + not has_moonep(),
43 + reason="Requires MoonEP",
44 +)
45 +
46 +
47 +class MulticastNotAvailableError(RuntimeError):
48 + pass
49 +
50 +
51 +@dataclasses.dataclass
52 +class TestConfig:
53 + topk: int
54 + m: int
55 + k: int
56 + n: int
57 + num_experts: int
58 + router_skew: float
59 + apply_router_weight_on_input: bool = False
60 +
61 +
62 +@dataclasses.dataclass
63 +class TestTensors:
64 + rank_tokens: torch.Tensor
65 + topk: torch.Tensor
66 + topk_weights: torch.Tensor
67 + config: TestConfig
68 +
69 + @staticmethod
70 + def make(config: TestConfig) -> "TestTensors":
71 + rank_tokens = (
72 + torch.randn((config.m, config.k), device="cuda", dtype=torch.bfloat16) / 10
73 + )
74 + # Skewed routing via a per-expert bias: scaling logits by a scalar
75 + # preserves their ordering and would not change the selected topk
76 + # ids, but a shared bias makes a few experts hot on every token so
77 + # the planner has to fill redundant-expert prefetch slots.
78 + expert_bias = config.router_skew * torch.randn(
79 + 1, config.num_experts, device="cuda", dtype=torch.float32
80 + )
81 + logits = (
82 + torch.randn(
83 + config.m, config.num_experts, device="cuda", dtype=torch.float32
84 + )
85 + + expert_bias
86 + )
87 + topk_weights, topk = torch.topk(logits, config.topk, dim=-1)
88 + topk_weights = torch.softmax(topk_weights, dim=-1)
89 + return TestTensors(
90 + rank_tokens=rank_tokens,
91 + topk=topk.to(dtype=torch.int64),
92 + topk_weights=topk_weights,
93 + config=config,
94 + )
95 +
96 +
97 +def reference_moonep_experts(
98 + hidden_nvsh: torch.Tensor,
99 + route_weights_nvs: torch.Tensor,
100 + cu_seqlens: torch.Tensor,
101 + weight_layout: "MoonEPExpertWeightLayout",
102 +) -> torch.Tensor:
103 + """Segment loop over MoonEP's expert-grouped layout.
104 +
105 + ``prefetch_weight`` has already materialized redundant experts' weights
106 + in rows ``[E, E+B)``, so every segment reads its own row. Route weights
107 + are applied here; MoonEP's combine does the K-sum.
108 + """
109 + output = torch.empty_like(hidden_nvsh)
110 + prev = 0
111 + for row, cur in enumerate(cu_seqlens.tolist()):
112 + if cur == prev:
113 + continue
114 + x = hidden_nvsh[prev:cur]
115 + gate = F.linear(x, weight_layout.full_gate_weight[row])
116 + up = F.linear(x, weight_layout.full_up_weight[row])
117 + y = F.linear(F.silu(gate) * up, weight_layout.full_down_weight[row])
118 + y = y * route_weights_nvs[prev:cur].to(dtype=y.dtype).unsqueeze(-1)
119 + output[prev:cur].copy_(y)
120 + prev = cur
121 + if prev < hidden_nvsh.shape[0]:
122 + output[prev:].zero_()
123 + return output
124 +
125 +
126 +def make_moonep_prepare_finalize(
127 + pg: ProcessGroup,
128 + pgi: ProcessGroupInfo,
129 + hidden_size: int,
130 + num_experts: int,
131 + topk: int,
132 + max_tokens_per_rank: int,
133 + weight_layout: "MoonEPExpertWeightLayout",
134 + pass_layout_to_pf: bool = True,
135 +):
136 + from moonep._C import nvl_multicast_supported
137 +
138 + from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
139 + MoonEPBufferPool,
140 + )
141 +
142 + if not nvl_multicast_supported():
143 + raise MulticastNotAvailableError("NVSwitch multicast not available")
144 +
145 + pool = MoonEPBufferPool(
146 + dict(
147 + H=hidden_size,
148 + K=topk,
149 + E=num_experts,
150 + num_ep_ranks=pgi.world_size,
151 + B=weight_layout.num_prefetch_slots,
152 + group=pg,
153 + explicitly_destroy=True,
154 + ),
155 + max_tokens_per_rank=max_tokens_per_rank,
156 + )
157 + return pool, MoonEPPrepareAndFinalize(
158 + buffer_pool=pool,
159 + max_tokens_per_rank=max_tokens_per_rank,
160 + num_dispatchers=pgi.world_size,
161 + num_global_experts=num_experts,
162 + weight_layout=weight_layout if pass_layout_to_pf else None,
163 + )
164 +
165 +
166 +def moonep_moe_impl(
167 + pg: ProcessGroup,
168 + pgi: ProcessGroupInfo,
169 + test_tensors: TestTensors,
170 + w1: torch.Tensor,
171 + w2: torch.Tensor,
172 + num_prefetch_slots: int,
173 +) -> torch.Tensor:
174 + config = test_tensors.config
175 + hidden_size = test_tensors.rank_tokens.shape[1]
176 + max_tokens_per_rank = 128 * ((config.m + 127) // 128)
177 +
178 + weight_layout = make_moonep_weight_layout(w1, w2, num_prefetch_slots)
179 + pool, pf = make_moonep_prepare_finalize(
180 + pg,
181 + pgi,
182 + hidden_size,
183 + config.num_experts,
184 + config.topk,
185 + max_tokens_per_rank,
186 + weight_layout,
187 + )
188 + try:
189 + hidden_nvsh, _, _, _, route_weights_nvs = pf.prepare(
190 + test_tensors.rank_tokens,
191 + test_tensors.topk_weights,
192 + test_tensors.topk,
193 + num_experts=config.num_experts,
194 + expert_map=None,
195 + apply_router_weight_on_input=config.apply_router_weight_on_input,
196 + quant_config=_no_quant_config(),
197 + )
198 + assert route_weights_nvs is not None
199 + if config.apply_router_weight_on_input:
200 + # prepare() already applied the weights to the inputs; the
201 + # runner must not apply them again.
202 + route_weights_nvs = torch.ones_like(route_weights_nvs)
203 + if config.router_skew >= 8 and config.m >= 100:
204 + # Heavy skew must actually engage the redundant-expert planner.
205 + assert (pf.plan.experts_to_copy >= 0).any(), (
206 + "no prefetch slot used despite heavy router skew"
207 + )
208 + expert_out = reference_moonep_experts(
209 + hidden_nvsh, route_weights_nvs, pf.cu_seqlens, weight_layout
210 + )
211 + output = torch.empty_like(test_tensors.rank_tokens)
212 + pf.finalize(
213 + output,
214 + expert_out,
215 + test_tensors.topk_weights,
216 + test_tensors.topk,
217 + apply_router_weight_on_input=False,
218 + weight_and_reduce_impl=TopKWeightAndReduceNoOP(),
219 + )
220 + torch.accelerator.synchronize()
221 + return output
222 + finally:
223 + pool.destroy()
224 +
225 +
226 +def _no_quant_config():
227 + from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
228 +
229 + return FusedMoEQuantConfig.make(quant_dtype=None)
230 +
231 +
232 +def moonep_modular_kernel_impl(
233 + pg: ProcessGroup,
234 + pgi: ProcessGroupInfo,
235 + test_tensors: TestTensors,
236 + w1: torch.Tensor,
237 + w2: torch.Tensor,
238 + num_prefetch_slots: int,
239 +) -> torch.Tensor:
240 + """Full modular-kernel path: MoonEPPrepareAndFinalize + MoonEPExperts."""
241 + from tests.kernels.moe.utils import make_dummy_moe_config
242 + from vllm.model_executor.layers.fused_moe.experts.moonep_experts import (
243 + MoonEPExperts,
244 + )
245 + from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
246 +
247 + config = test_tensors.config
248 + hidden_size = test_tensors.rank_tokens.shape[1]
249 + max_tokens_per_rank = 128 * ((config.m + 127) // 128)
250 +
251 + weight_layout = make_moonep_weight_layout(w1, w2, num_prefetch_slots)
252 + # weight_layout deliberately not passed to the P/F: the engine path
253 + # resolves it through post_init_setup + the experts' weight hook.
254 + pool, pf = make_moonep_prepare_finalize(
255 + pg,
256 + pgi,
257 + hidden_size,
258 + config.num_experts,
259 + config.topk,
260 + max_tokens_per_rank,
261 + weight_layout,
262 + pass_layout_to_pf=False,
263 + )
264 + try:
265 + moe_config = make_dummy_moe_config(
266 + num_experts=config.num_experts,
267 + experts_per_token=config.topk,
268 + hidden_dim=hidden_size,
269 + intermediate_size=config.n,
270 + max_num_tokens=max_tokens_per_rank,
271 + )
272 + experts = MoonEPExperts(moe_config=moe_config, quant_config=_no_quant_config())
273 + # Stand-in for the layer carrying the layout, as
274 + # convert_to_unquantized_kernel_format leaves it in the engine path.
275 + fake_layer = types.SimpleNamespace(_moonep_weight_layout=weight_layout)
276 + experts.process_weights_after_loading(fake_layer)
277 + kernel = FusedMoEKernel(prepare_finalize=pf, fused_experts=experts)
278 + out = kernel.apply(
279 + hidden_states=test_tensors.rank_tokens,
280 + w1=weight_layout.full_gate_weight,
281 + w2=weight_layout.full_down_weight,
282 + topk_weights=test_tensors.topk_weights,
283 + topk_ids=test_tensors.topk,
284 + activation=MoEActivation.SILU,
285 + global_num_experts=config.num_experts,
286 + expert_map=None,
287 + apply_router_weight_on_input=config.apply_router_weight_on_input,
288 + )
289 + torch.accelerator.synchronize()
290 + return out
291 + finally:
292 + pool.destroy()
293 +
294 +
295 +def _moonep_moe(
296 + pgi: ProcessGroupInfo,
297 + config: TestConfig,
298 + w1: torch.Tensor,
299 + w2: torch.Tensor,
300 + num_prefetch_slots: int,
301 + use_modular_kernel: bool,
302 +):
303 + from vllm.v1.worker.workspace import init_workspace_manager
304 +
305 + device_idx = torch.accelerator.current_device_index()
306 + init_workspace_manager(torch.device("cuda", device_idx))
307 + w1 = w1.to(device=device_idx)
308 + w2 = w2.to(device=device_idx)
309 +
310 + pg = torch.distributed.new_group(list(range(pgi.world_size)))
311 + set_random_seed(7 + pgi.rank)
312 + test_tensors = TestTensors.make(config)
313 +
314 + with set_current_vllm_config(VllmConfig()):
315 + torch_combined = torch_experts(
316 + test_tensors.rank_tokens,
317 + w1,
318 + w2,
319 + test_tensors.topk_weights,
320 + test_tensors.topk,
321 + apply_router_weights_on_input=config.apply_router_weight_on_input,
322 + )
323 + impl = moonep_modular_kernel_impl if use_modular_kernel else moonep_moe_impl
324 + moonep_combined = impl(pg, pgi, test_tensors, w1, w2, num_prefetch_slots)
325 +
326 + torch.testing.assert_close(
327 + torch_combined,
328 + moonep_combined,
329 + atol=6e-2,
330 + rtol=6e-2,
331 + )
332 +
333 +
334 +MNKs = [
335 + (1, 256, 512),
336 + (37, 256, 512),
337 + (100, 512, 1024),
338 + (512, 768, 2048),
339 +]
340 +
341 +
342 +@pytest.mark.parametrize("m,n,k", MNKs)
343 +@pytest.mark.parametrize("num_experts", [32])
344 +@pytest.mark.parametrize("topk", [4])
345 +@pytest.mark.parametrize("router_skew", [1.0, 8.0])
346 +@pytest.mark.parametrize("num_prefetch_slots", [4])
347 +@pytest.mark.parametrize("world_size", [2])
348 +@pytest.mark.parametrize("use_modular_kernel", [False, True])
349 +@multi_gpu_test(num_gpus=2)
350 +@requires_moonep
351 +def test_moonep_bf16_moe(
352 + m: int,
353 + n: int,
354 + k: int,
355 + num_experts: int,
356 + topk: int,
357 + router_skew: float,
358 + num_prefetch_slots: int,
359 + world_size: int,
360 + use_modular_kernel: bool,
361 +):
362 + set_random_seed(7)
363 + config = TestConfig(
364 + topk=topk,
365 + m=m,
366 + k=k,
367 + n=n,
368 + num_experts=num_experts,
369 + router_skew=router_skew,
370 + )
371 + (_, w1, _, _), (_, w2, _, _) = make_test_weights(num_experts, n, k)
372 +
373 + try:
374 + parallel_launch(
375 + world_size,
376 + _moonep_moe,
377 + config,
378 + w1,
379 + w2,
380 + num_prefetch_slots,
381 + use_modular_kernel,
382 + )
383 + except Exception as exc:
384 + if "MulticastNotAvailableError" in str(exc):
385 + pytest.skip("NVSwitch multicast not available")
386 + raise
387 +
388 +
389 +@pytest.mark.parametrize("use_modular_kernel", [False, True])
390 +@multi_gpu_test(num_gpus=2)
391 +@requires_moonep
392 +def test_moonep_input_weighted_moe(use_modular_kernel: bool):
393 + """Llama-4-style routing applies route weights to the expert inputs."""
394 + set_random_seed(7)
395 + config = TestConfig(
396 + topk=1,
397 + m=100,
398 + k=1024,
399 + n=512,
400 + num_experts=32,
401 + router_skew=1.0,
402 + apply_router_weight_on_input=True,
403 + )
404 + (_, w1, _, _), (_, w2, _, _) = make_test_weights(
405 + config.num_experts, config.n, config.k
406 + )
407 + try:
408 + parallel_launch(2, _moonep_moe, config, w1, w2, 4, use_modular_kernel)
409 + except Exception as exc:
410 + if "MulticastNotAvailableError" in str(exc):
411 + pytest.skip("NVSwitch multicast not available")
412 + raise
vllm/config/parallel.py
+2
@@ -49,6 +49,7 @@ All2AllBackend = Literal[
49 "deepep_v2",
50 "mori_high_throughput",
51 "mori_low_latency",
52 + "moonep",
53 "nixl_ep",
54 "allgather_reducescatter",
55 "flashinfer_all2allv", # temporary alias for flashinfer_nvlink_two_sided
@@ -206,6 +207,7 @@ class ParallelConfig:
207 - "deepep_low_latency": Use deepep low-latency kernels
208 - "mori_high_throughput": MoRI EP with InterNodeV1 for multi-node
209 - "mori_low_latency": MoRI EP with InterNodeV1LL for multi-node
210 + - "moonep": MoonEP balanced EP with dynamic redundant experts (NVLink)
211 - "nixl_ep": Use nixl-ep kernels
212 - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels
213 - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl"""
vllm/config/vllm.py
+49
@@ -1881,6 +1881,55 @@ class VllmConfig:
1881 # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled)
1882 self._set_compile_ranges()
1883
1884 + if self.parallel_config.all2all_backend == "moonep":
1885 + if (
1886 + self.model_config is not None
1887 + and self.model_config.quantization is not None
1888 + ):
1889 + raise ValueError(
1890 + "The moonep all2all backend currently supports unquantized "
1891 + "BF16 models only; got "
1892 + f"quantization={self.model_config.quantization!r}. Use a "
1893 + "different --all2all-backend for quantized models."
1894 + )
1895 + if (
1896 + self.model_config is not None
1897 + and self.model_config.dtype != torch.bfloat16
1898 + ):
1899 + raise ValueError(
1900 + "The moonep all2all backend currently supports BF16 models "
1901 + f"only; got dtype={self.model_config.dtype}. Use a "
1902 + "different --all2all-backend or --dtype bfloat16."
1903 + )
1904 + if self.parallel_config.enable_eplb:
1905 + raise ValueError(
1906 + "The moonep all2all backend does not support EPLB yet: "
1907 + "EPLB rearranges expert parameters in a layout MoonEP's "
1908 + "replicated [E+B] weights do not follow. Disable "
1909 + "--enable-eplb or use a different --all2all-backend."
1910 + )
1911 + if self.parallel_config.expert_placement_strategy != "linear":
1912 + raise ValueError(
1913 + "The moonep all2all backend requires linear expert "
1914 + "placement: its load-time all-gather assumes each rank "
1915 + "holds a contiguous chunk of the global expert range. Got "
1916 + "--expert-placement-strategy "
1917 + f"{self.parallel_config.expert_placement_strategy!r}."
1918 + )
1919 + # Enforced here rather than in set_splitting_ops_for_v1 so it
1920 + # holds for every compilation mode, and keyed on use_all2all so
1921 + # PCP/SP-only topologies are covered too: MoonEP dispatch/combine
1922 + # are eager-only and must not be captured.
1923 + if (
1924 + self.parallel_config.use_all2all
1925 + and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
1926 + ):
1927 + logger.info(
1928 + "MoonEP: Disabling CUDA Graphs since the MoonEP "
1929 + "integration is currently eager-only."
1930 + )
1931 + self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
1932 +
1933 # Do this after all the updates to compilation_config.mode
1934 effective_dp_size = (
1935 self.parallel_config.data_parallel_size
vllm/distributed/device_communicators/all2all.py
+74 -1
@@ -21,7 +21,12 @@ from vllm.utils.flashinfer import (
21 has_flashinfer_nvlink_two_sided,
22 )
23 from vllm.utils.func_utils import supports_kw
24 -from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori
24 +from vllm.utils.import_utils import (
25 + check_moonep_system_support,
26 + has_deep_ep,
27 + has_deep_ep_v2,
28 + has_mori,
29 +)
30
31 from .base_device_communicator import All2AllManagerBase, Cache
32
@@ -1094,3 +1099,71 @@ class DeepEPV2All2AllManager(All2AllManagerBase):
1099 for _, handle in self.handle_cache._cache.items():
1100 handle.destroy()
1101 self.handle_cache._cache.clear()
1102 +
1103 +
1104 +class MoonEPAll2AllManager(All2AllManagerBase):
1105 + """All2All communication based on MoonEP
1106 + (https://github.com/MoonshotAI/MoonEP).
1107 +
1108 + MoonEP keeps token loads perfectly balanced across EP ranks by planning
1109 + a small number of dynamically redundant experts online and prefetching
1110 + their weights before expert compute. Every rank receives exactly
1111 + S x K token slots regardless of router skew, so all communication and
1112 + compute shapes are static.
1113 +
1114 + Requires NVLink symmetric-memory / multicast capable topologies
1115 + (single node NVSwitch, e.g. H100/H200/B200/GB300 class).
1116 + """
1117 +
1118 + def __init__(self, cpu_group, tcp_store_group=None, device_group=None):
1119 + check_moonep_system_support()
1120 + super().__init__(cpu_group, tcp_store_group)
1121 + self._device_group = device_group
1122 + self.handle_cache = Cache()
1123 +
1124 + def _make_buffer_kwargs(
1125 + self,
1126 + max_num_tokens_per_dp_rank: int,
1127 + token_hidden_size: int,
1128 + num_topk: int,
1129 + num_global_experts: int,
1130 + num_prefetch_slots: int,
1131 + token_padding: int,
1132 + num_sms: int,
1133 + ) -> dict:
1134 + return dict(
1135 + S=max_num_tokens_per_dp_rank,
1136 + H=token_hidden_size,
1137 + K=num_topk,
1138 + E=num_global_experts,
1139 + num_ep_ranks=self.world_size,
1140 + num_sms=num_sms,
1141 + token_padding=token_padding,
1142 + B=num_prefetch_slots,
1143 + group=self._device_group
1144 + if self._device_group is not None
1145 + else self.cpu_group,
1146 + explicitly_destroy=True,
1147 + )
1148 +
1149 + def get_handle(self, kwargs):
1150 + from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
1151 + MoonEPBufferPool,
1152 + )
1153 +
1154 + buffer_kwargs = self._make_buffer_kwargs(**kwargs)
1155 + logger.debug("MoonEP all2all args %s", buffer_kwargs)
1156 +
1157 + def make_pool(**kw):
1158 + return MoonEPBufferPool(kw, max_tokens_per_rank=kw["S"])
1159 +
1160 + handle: MoonEPBufferPool = self.handle_cache.get_or_create(
1161 + buffer_kwargs, make_pool
1162 + )
1163 + return handle
1164 +
1165 + def destroy(self):
1166 + with self.handle_cache._lock:
1167 + for _, handle in self.handle_cache._cache.items():
1168 + handle.destroy()
1169 + self.handle_cache._cache.clear()
vllm/distributed/device_communicators/cuda_communicator.py
+8
@@ -214,6 +214,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
214 tcp_store_group,
215 device_group=self.device_group,
216 )
217 + elif self.all2all_backend == "moonep":
218 + from .all2all import MoonEPAll2AllManager
219 +
220 + self.all2all_manager = MoonEPAll2AllManager(
221 + self.cpu_group,
222 + tcp_store_group,
223 + device_group=self.device_group,
224 + )
225 elif self.all2all_backend == "nixl_ep":
226 from .all2all import NixlEPAll2AllManager
227
vllm/model_executor/layers/fused_moe/all2all_utils.py
+30
@@ -35,6 +35,7 @@ from vllm.platforms import current_platform
35 from vllm.utils.import_utils import (
36 has_deep_ep,
37 has_deep_ep_v2,
38 + has_moonep,
39 has_mori,
40 has_nixl_ep,
41 )
@@ -103,6 +104,13 @@ if current_platform.is_cuda_alike():
104 NIXL_EP_QUANT_BLOCK_SHAPE,
105 NixlEPPrepareAndFinalize,
106 )
107 + if has_moonep():
108 + from .prepare_finalize.moonep import (
109 + MOONEP_DEFAULT_NUM_PREFETCH_SLOTS,
110 + MOONEP_DEFAULT_NUM_SMS,
111 + MOONEP_DEFAULT_TOKEN_PADDING,
112 + MoonEPPrepareAndFinalize,
113 + )
114
115
116 def get_ep_all2all_manager(
@@ -279,6 +287,28 @@ def maybe_make_prepare_finalize(
287 sp_size=moe.moe_parallel_config.sp_size,
288 )
289
290 + elif moe.use_moonep_kernels:
291 + all_to_all_args = dict(
292 + max_num_tokens_per_dp_rank=moe.max_num_tokens,
293 + token_hidden_size=moe.hidden_dim,
294 + num_topk=moe.experts_per_token,
295 + num_global_experts=moe.num_experts,
296 + num_prefetch_slots=MOONEP_DEFAULT_NUM_PREFETCH_SLOTS,
297 + token_padding=MOONEP_DEFAULT_TOKEN_PADDING,
298 + num_sms=MOONEP_DEFAULT_NUM_SMS,
299 + )
300 + handle = all2all_manager.get_handle(all_to_all_args)
301 +
302 + # The [E+B] weight layout is picked up from the experts (their
303 + # process_weights_after_loading hook) once the layer has loaded and
304 + # converted its weights.
305 + prepare_finalize = MoonEPPrepareAndFinalize(
306 + buffer_pool=handle,
307 + max_tokens_per_rank=moe.max_num_tokens,
308 + num_dispatchers=all2all_manager.world_size,
309 + num_global_experts=moe.num_experts,
310 + )
311 +
312 elif moe.use_mori_kernels:
313 assert quant_config is not None
314
vllm/model_executor/layers/fused_moe/config.py
+8
@@ -1085,6 +1085,10 @@ class FusedMoEParallelConfig:
1085 def use_deepep_v2_kernels(self):
1086 return self.use_all2all_kernels and self.all2all_backend == "deepep_v2"
1087
1088 + @property
1089 + def use_moonep_kernels(self):
1090 + return self.use_all2all_kernels and self.all2all_backend == "moonep"
1091 +
1092 @staticmethod
1093 def flatten_tp_across_dp_and_pcp(
1094 tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
@@ -1473,6 +1477,10 @@ class FusedMoEConfig:
1477 def use_deepep_v2_kernels(self):
1478 return self.moe_parallel_config.use_deepep_v2_kernels
1479
1480 + @property
1481 + def use_moonep_kernels(self):
1482 + return self.moe_parallel_config.use_moonep_kernels
1483 +
1484 @property
1485 def needs_round_robin_routing_tables(self):
1486 return self.moe_parallel_config.needs_round_robin_routing_tables
vllm/model_executor/layers/fused_moe/experts/moonep_experts.py new
+220
@@ -0,0 +1,220 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MoonEP experts: grouped GEMM over MoonEP's expert-grouped activations.
4 +
5 +MoonEP dispatch delivers tokens already contiguous per weight row
6 +(``[NvS, H]`` in ``cu_seqlens[E+B]`` segment order), so the expert compute
7 +is three grouped GEMMs over those segments with no permute/unpermute:
8 +
9 + gate = grouped_mm(x, w_gate[E+B]) # [NvS, I]
10 + up = grouped_mm(x, w_up[E+B]) # [NvS, I]
11 + act = silu(gate) * up * route_weight # route weights applied here
12 + out = grouped_mm(act, w_down[E+B]) # [NvS, H]
13 +
14 +Rows ``[E, E+B)`` are the redundant-expert prefetch slots that
15 +``MoonEPPrepareAndFinalize`` fills before ``apply`` runs. Empty segments
16 +(including unused prefetch slots) are skipped by the grouped GEMM.
17 +
18 +Weight layout: MoonEP's ``prefetch_weight`` requires each projection to be
19 +its own contiguous ``[E+B, ., .]`` tensor, so gate and up are separate
20 +tensors rather than vLLM's interleaved ``[E, 2I, H]`` ``w13``. The modular
21 +kernel passes ``w1``/``w2`` from the layer (``w1`` is the gate tensor); the
22 +full layout is picked up from the layer in ``process_weights_after_loading``
23 +and shared with ``MoonEPPrepareAndFinalize`` via ``post_init_setup``.
24 +"""
25 +
26 +import torch
27 +
28 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
29 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
30 +from vllm.model_executor.layers.fused_moe.config import (
31 + FusedMoEConfig,
32 + FusedMoEParallelConfig,
33 + FusedMoEQuantConfig,
34 +)
35 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
36 + TopKWeightAndReduceNoOP,
37 +)
38 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
39 +from vllm.platforms import current_platform
40 +
41 +
42 +def moonep_grouped_gemm(
43 + a: torch.Tensor,
44 + w: torch.Tensor,
45 + cu_seqlens: torch.Tensor,
46 +) -> torch.Tensor:
47 + """``out[seg_g] = a[seg_g] @ w[g].T`` for each ``cu_seqlens`` segment.
48 +
49 + Args:
50 + a: ``[NvS, K]`` activations in segment order.
51 + w: ``[G, N, K]`` weights (vLLM ``[E, N, K]`` convention).
52 + cu_seqlens: ``[G]`` int32 cumulative segment end offsets.
53 +
54 + Returns:
55 + ``[NvS, N]``; rows past ``cu_seqlens[-1]`` are zero-filled.
56 +
57 + """
58 + return torch._grouped_mm(a, w.transpose(1, 2), offs=cu_seqlens)
59 +
60 +
61 +class MoonEPExperts(mk.FusedMoEExpertsModular):
62 + """Grouped-GEMM experts over MoonEP's ``[NvS, H]`` layout (BF16)."""
63 +
64 + def __init__(
65 + self,
66 + moe_config: FusedMoEConfig,
67 + quant_config: FusedMoEQuantConfig,
68 + max_num_tokens: int | None = None,
69 + num_dispatchers: int | None = None,
70 + ):
71 + super().__init__(moe_config, quant_config, max_num_tokens, num_dispatchers)
72 + assert quant_config.quant_dtype is None, (
73 + "MoonEPExperts supports unquantized BF16 only"
74 + )
75 + if (
76 + quant_config.gemm1_alpha is not None
77 + or quant_config.gemm1_beta is not None
78 + or quant_config.gemm1_clamp_limit is not None
79 + ):
80 + raise NotImplementedError(
81 + "MoonEPExperts implements plain silu(gate) * up only; SwiGLU "
82 + "variants with alpha/beta/clamp parameters are not supported"
83 + )
84 + self._weight_layout = None
85 +
86 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
87 + # The [E+B] gate/up/down layout built by
88 + # convert_to_unquantized_kernel_format; MoonEPPrepareAndFinalize
89 + # reads it from here (post_init_setup) for weight prefetch.
90 + layout = getattr(layer, "_moonep_weight_layout", None)
91 + assert layout is not None, (
92 + "MoonEPExperts requires the layer to carry _moonep_weight_layout "
93 + "(set by convert_to_unquantized_kernel_format)"
94 + )
95 + self._weight_layout = layout
96 +
97 + @property
98 + def weight_layout(self):
99 + return self._weight_layout
100 +
101 + @staticmethod
102 + def activation_format() -> mk.FusedMoEActivationFormat:
103 + return mk.FusedMoEActivationFormat.Standard
104 +
105 + @staticmethod
106 + def _supports_current_device() -> bool:
107 + return current_platform.is_cuda()
108 +
109 + @staticmethod
110 + def _supports_no_act_and_mul() -> bool:
111 + return False
112 +
113 + @staticmethod
114 + def _supports_activation(activation: MoEActivation) -> bool:
115 + return activation == MoEActivation.SILU
116 +
117 + @staticmethod
118 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
119 + # EPLB rearranges the named expert parameters as local_num_experts
120 + # rows, which does not understand the replicated [E+B] layout.
121 + return (
122 + moe_parallel_config.use_moonep_kernels
123 + and not moe_parallel_config.enable_eplb
124 + )
125 +
126 + @staticmethod
127 + def _supports_quant_scheme(
128 + weight_key: QuantKey | None,
129 + activation_key: QuantKey | None,
130 + ) -> bool:
131 + return weight_key is None and activation_key is None
132 +
133 + def supports_chunking(self) -> bool:
134 + # NvS is static and segments must stay whole.
135 + return False
136 +
137 + def supports_expert_map(self) -> bool:
138 + # MoonEP addresses global expert rows directly.
139 + return False
140 +
141 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
142 + return TopKWeightAndReduceNoOP()
143 +
144 + def moe_problem_size(
145 + self,
146 + a1: torch.Tensor,
147 + w1: torch.Tensor,
148 + w2: torch.Tensor,
149 + topk_ids: torch.Tensor,
150 + ) -> tuple[int, int, int, int, int]:
151 + # a1 is [NvS, H] in segment order, not token-major, so the base
152 + # implementation's topk_ids/a1 row-count check does not apply.
153 + assert a1.dim() == 2 and w1.dim() == 3 and w2.dim() == 3
154 + num_rows, intermediate, hidden = w1.shape # [E+B, I, H]
155 + assert a1.size(1) == hidden
156 + return num_rows, a1.size(0), intermediate, hidden, topk_ids.size(1)
157 +
158 + def workspace_shapes(
159 + self,
160 + M: int,
161 + N: int,
162 + K: int,
163 + topk: int,
164 + global_num_experts: int,
165 + local_num_experts: int,
166 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
167 + activation: MoEActivation,
168 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
169 + # M is NvS. torch._grouped_mm allocates its own outputs, so no
170 + # workspaces are needed.
171 + return ((0,), (0,), (M, K))
172 +
173 + def apply(
174 + self,
175 + output: torch.Tensor,
176 + hidden_states: torch.Tensor,
177 + w1: torch.Tensor,
178 + w2: torch.Tensor,
179 + topk_weights: torch.Tensor,
180 + topk_ids: torch.Tensor,
181 + activation: MoEActivation,
182 + global_num_experts: int,
183 + expert_map: torch.Tensor | None,
184 + a1q_scale: torch.Tensor | None,
185 + a2_scale: torch.Tensor | None,
186 + workspace13: torch.Tensor,
187 + workspace2: torch.Tensor,
188 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
189 + apply_router_weight_on_input: bool,
190 + ) -> None:
191 + # expert_map is ignored: MoonEP dispatches on global expert ids and
192 + # the weights are addressed by global [E+B] row, on every rank.
193 + assert activation == MoEActivation.SILU
194 + assert self._weight_layout is not None, (
195 + "process_weights_after_loading() not called"
196 + )
197 + assert expert_tokens_meta is not None
198 + # MoonEPPrepareAndFinalize returns route weights in NvS order as the
199 + # dispatched topk_weights, and per-row segment lengths as
200 + # expert_tokens_meta.
201 + route_weights_nvs = topk_weights
202 + assert route_weights_nvs.dim() == 1
203 + cu_seqlens = torch.cumsum(
204 + expert_tokens_meta.expert_num_tokens, dim=0, dtype=torch.int32
205 + )
206 +
207 + assert cu_seqlens.numel() == w1.size(0)
208 +
209 + gate = moonep_grouped_gemm(hidden_states, w1, cu_seqlens)
210 + up = moonep_grouped_gemm(
211 + hidden_states, self._weight_layout.full_up_weight, cu_seqlens
212 + )
213 + act = torch.nn.functional.silu(gate)
214 + act.mul_(up)
215 + if not apply_router_weight_on_input:
216 + # Input-weighted routing (top-1, e.g. Llama 4) was already
217 + # applied to the activations by prepare(), before dispatch.
218 + act.mul_(route_weights_nvs.to(act.dtype).unsqueeze(-1))
219 + # Padding rows past the last segment come out zero-filled.
220 + output.copy_(moonep_grouped_gemm(act, w2, cu_seqlens))
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+52 -1
@@ -40,6 +40,7 @@ class UnquantizedMoeBackend(Enum):
40 AITER = "ROCm AITER"
41 TRITON = "TRITON"
42 BATCHED_TRITON = "BATCHED_TRITON"
43 + MOONEP = "MOONEP"
44 CPU = "CPU"
45 XPU = "XPU"
46 TPU = "TPU"
@@ -138,6 +139,13 @@ def backend_to_kernel_cls(
139
140 return [BatchedTritonExperts]
141
142 + elif backend == UnquantizedMoeBackend.MOONEP:
143 + from vllm.model_executor.layers.fused_moe.experts.moonep_experts import (
144 + MoonEPExperts,
145 + )
146 +
147 + return [MoonEPExperts]
148 +
149 elif backend == UnquantizedMoeBackend.XPU:
150 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExperts
151
@@ -281,6 +289,13 @@ def select_unquantized_moe_backend(
289 return backend, k_cls
290 raise ValueError(_make_log_unsupported(backend, reason))
291
292 + # MoonEP owns the expert layout (expert-grouped [NvS, H] + cu_seqlens),
293 + # so it is not interchangeable with the token-major experts backends.
294 + if moe_config.moe_parallel_config.use_moonep_kernels:
295 + return _return_or_raise(
296 + UnquantizedMoeBackend.MOONEP, moe_config, activation_format
297 + )
298 +
299 runner_backend = moe_config.moe_backend
300 # 'humming' is quantization-only; an unquantized layer (e.g. excluded via
301 # modules_to_not_convert) falls through to auto instead of erroring.
@@ -366,7 +381,42 @@ def convert_to_unquantized_kernel_format(
381 moe_config: FusedMoEConfig,
382 w13_weight: torch.Tensor,
383 w2_weight: torch.Tensor,
384 + layer: torch.nn.Module | None = None,
385 ) -> tuple[torch.Tensor, torch.Tensor]:
386 + if unquantized_backend == UnquantizedMoeBackend.MOONEP:
387 + if layer is None:
388 + raise ValueError(
389 + "MoonEP weight conversion requires the layer module to stash "
390 + "the [E+B] weight layout on"
391 + )
392 + # MoonEP addresses expert weights by global [E+B] row and needs one
393 + # contiguous tensor per projection (gate / up / down). The layer's
394 + # w13_weight becomes the gate tensor and w2_weight the down tensor;
395 + # the up tensor and the full layout are stashed on the layer for
396 + # the kernel wiring in _setup_kernel.
397 + from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
398 + MOONEP_DEFAULT_NUM_PREFETCH_SLOTS,
399 + gather_moonep_weight_layout,
400 + )
401 +
402 + if getattr(layer, "_moonep_weight_layout", None) is not None:
403 + # On a reload the layer parameters were already replaced by the
404 + # converted [E+B] gate/down tensors, so the source-format weights
405 + # this conversion needs are gone (and the unregistered up
406 + # projection cannot be recovered from them at all). Tracked in
407 + # RFC #52095.
408 + raise NotImplementedError(
409 + "MoonEP does not support in-place weight reloads yet"
410 + )
411 + layout = gather_moonep_weight_layout(
412 + w13_weight,
413 + w2_weight,
414 + num_global_experts=moe_config.num_experts,
415 + num_prefetch_slots=MOONEP_DEFAULT_NUM_PREFETCH_SLOTS,
416 + )
417 + layer._moonep_weight_layout = layout
418 + return layout.full_gate_weight, layout.full_down_weight
419 +
420 if unquantized_backend == UnquantizedMoeBackend.AITER:
421 w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight)
422 w13_weight.is_shuffled = True
@@ -522,9 +572,10 @@ class UnquantizedMoEKernelOracle(MoEKernelOracle[UnquantizedMoeBackend]):
572 moe_config: FusedMoEConfig,
573 w13_weight: torch.Tensor,
574 w2_weight: torch.Tensor,
575 + layer: torch.nn.Module | None = None,
576 ) -> tuple[torch.Tensor, torch.Tensor]:
577 return convert_to_unquantized_kernel_format(
527 - backend, moe_config, w13_weight, w2_weight
578 + backend, moe_config, w13_weight, w2_weight, layer=layer
579 )
580
581 def make_kernel(
vllm/model_executor/layers/fused_moe/prepare_finalize/moonep.py new
+429
@@ -0,0 +1,429 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MoonEP (https://github.com/MoonshotAI/MoonEP) prepare/finalize.
4 +
5 +BF16 correctness-first proof of concept on top of vLLM's modular kernel
6 +interface.
7 +
8 +MoonEP differs from DeepEP-style backends in two ways that shape this
9 +integration:
10 +
11 +- ``dispatch`` returns tokens already grouped by expert *row* (a fixed
12 + ``[NvS, H]`` layout with ``NvS = S x K`` real slots plus padding) together
13 + with a ``cu_seqlens[E+B]`` segment table and an opaque ``plan``. There is
14 + no per-token topk id tensor after dispatch; the expert compute must be a
15 + grouped GEMM over ``cu_seqlens`` segments.
16 +- Rows ``[E, E+B)`` of the weight/segment space are dynamic redundant-expert
17 + prefetch slots. ``plan.experts_to_copy`` names the source expert of each
18 + slot and ``Buffer.prefetch_weight`` must run between dispatch and expert
19 + compute.
20 +
21 +PoC limitations:
22 +- BF16 / unquantized only, eager only.
23 +- Expert weights are replicated in global-expert order on every rank
24 + (memory-heavy). Production Kimi-K3 serving requires sharded
25 + symmetric-memory expert ownership, where rows ``[0, E)`` physically alias
26 + each home rank's parameter memory.
27 +- Route weights are applied inside the expert compute and MoonEP's
28 + ``combine`` performs the K-sum, so ``finalize`` requires
29 + ``TopKWeightAndReduceNoOP``.
30 +"""
31 +
32 +from typing import Any, NamedTuple
33 +
34 +import torch
35 +import torch.nn.functional as F
36 +
37 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
38 +from vllm.forward_context import get_forward_context
39 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
40 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
41 + TopKWeightAndReduceNoOP,
42 +)
43 +
44 +MOONEP_DEFAULT_NUM_PREFETCH_SLOTS = 4
45 +MOONEP_DEFAULT_TOKEN_PADDING = 128
46 +MOONEP_DEFAULT_NUM_SMS = 32
47 +# MoonEP's weight-prefetch kernel tiles hidden and intermediate dims by 128.
48 +MOONEP_WEIGHT_TILE = 128
49 +
50 +
51 +class MoonEPExpertWeightLayout(NamedTuple):
52 + """Contiguous BF16 expert weights in MoonEP ``[E+B, ...]`` layout.
53 +
54 + Rows ``[0, E)`` hold expert weights in global expert order; rows
55 + ``[E, E+B)`` are mutable prefetch slots filled by
56 + ``Buffer.prefetch_weight``.
57 +
58 + The prefetch slots never need re-zeroing between calls: the planner
59 + marks unused slots with ``-1`` in ``plan.experts_to_copy`` (skipped by
60 + ``prefetch_weight``) and gives them an empty ``cu_seqlens`` segment, so
61 + stale slot contents are never read; used slots are fully overwritten.
62 + """
63 +
64 + full_gate_weight: torch.Tensor
65 + full_up_weight: torch.Tensor
66 + full_down_weight: torch.Tensor
67 + num_prefetch_slots: int
68 +
69 +
70 +def make_moonep_weight_layout(
71 + w13_weight: torch.Tensor,
72 + w2_weight: torch.Tensor,
73 + num_prefetch_slots: int,
74 +) -> MoonEPExpertWeightLayout:
75 + """Build the replicated ``[E+B, ...]`` PoC weight layout.
76 +
77 + ``w13_weight`` must be ``[E, 2I, H]`` (gate rows first) and ``w2_weight``
78 + ``[E, H, I]``, both BF16 in global expert order.
79 + """
80 + if num_prefetch_slots <= 0:
81 + raise ValueError(
82 + f"num_prefetch_slots must be positive, got {num_prefetch_slots}"
83 + )
84 + if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16:
85 + raise NotImplementedError("MoonEP PoC supports BF16 weights only.")
86 + num_experts, two_i, hidden_size = w13_weight.shape
87 + intermediate_size = two_i // 2
88 + if tuple(w2_weight.shape) != (num_experts, hidden_size, intermediate_size):
89 + raise ValueError(
90 + f"w2_weight shape {tuple(w2_weight.shape)} does not match "
91 + f"w13_weight shape {tuple(w13_weight.shape)}"
92 + )
93 + if hidden_size % MOONEP_WEIGHT_TILE or intermediate_size % MOONEP_WEIGHT_TILE:
94 + raise ValueError(
95 + "MoonEP weight prefetch requires hidden_size and intermediate_size "
96 + f"to be multiples of {MOONEP_WEIGHT_TILE}; got H={hidden_size}, "
97 + f"I={intermediate_size}"
98 + )
99 +
100 + full_gate_weight = torch.empty(
101 + num_experts + num_prefetch_slots,
102 + intermediate_size,
103 + hidden_size,
104 + dtype=torch.bfloat16,
105 + device=w13_weight.device,
106 + )
107 + full_up_weight = torch.empty_like(full_gate_weight)
108 + full_down_weight = torch.empty(
109 + num_experts + num_prefetch_slots,
110 + hidden_size,
111 + intermediate_size,
112 + dtype=torch.bfloat16,
113 + device=w2_weight.device,
114 + )
115 +
116 + full_gate_weight[:num_experts].copy_(w13_weight[:, :intermediate_size, :])
117 + full_up_weight[:num_experts].copy_(w13_weight[:, intermediate_size:, :])
118 + full_down_weight[:num_experts].copy_(w2_weight)
119 + full_gate_weight[num_experts:].zero_()
120 + full_up_weight[num_experts:].zero_()
121 + full_down_weight[num_experts:].zero_()
122 +
123 + return MoonEPExpertWeightLayout(
124 + full_gate_weight=full_gate_weight.contiguous(),
125 + full_up_weight=full_up_weight.contiguous(),
126 + full_down_weight=full_down_weight.contiguous(),
127 + num_prefetch_slots=num_prefetch_slots,
128 + )
129 +
130 +
131 +def gather_moonep_weight_layout(
132 + w13_local: torch.Tensor,
133 + w2_local: torch.Tensor,
134 + num_global_experts: int,
135 + num_prefetch_slots: int,
136 +) -> MoonEPExpertWeightLayout:
137 + """Build the replicated ``[E+B, ...]`` layout from this rank's local experts.
138 +
139 + PoC bridge: each EP rank loads only its own experts (linear placement),
140 + so all-gather them once at load time into global expert order on every
141 + rank. Production MoonEP instead maps rows ``[0, E)`` onto each home
142 + rank's parameter memory via symmetric memory (RFC #52095 item 5).
143 + """
144 + from vllm.distributed import get_ep_group
145 +
146 + ep_group = get_ep_group()
147 + ep_size = ep_group.world_size
148 + if ep_size == 1:
149 + return make_moonep_weight_layout(w13_local, w2_local, num_prefetch_slots)
150 + if w13_local.size(0) * ep_size != num_global_experts:
151 + raise NotImplementedError(
152 + "MoonEP PoC requires num_experts to be evenly divisible across EP "
153 + f"ranks: {num_global_experts} experts, {ep_size} ranks"
154 + )
155 + w13_global = ep_group.all_gather(w13_local.contiguous(), dim=0)
156 + w2_global = ep_group.all_gather(w2_local.contiguous(), dim=0)
157 + return make_moonep_weight_layout(w13_global, w2_global, num_prefetch_slots)
158 +
159 +
160 +MOONEP_MIN_DISPATCH_TOKENS = 128
161 +
162 +
163 +class MoonEPBufferPool:
164 + """Lazily creates ``moonep.Buffer`` instances at power-of-two capacities.
165 +
166 + ``Buffer`` fixes its token capacity ``S`` at construction, so a single
167 + full-capacity buffer would pad every dispatch to
168 + ``max_num_batched_tokens``. The pool instead serves the smallest
169 + power-of-two capacity that fits the step's token count, so decode-sized
170 + batches dispatch a few hundred slots rather than the full capacity.
171 +
172 + Selection must be identical on every EP rank (dispatch is collective and
173 + ``Buffer`` construction is itself a collective): callers derive the token
174 + count from ``dp_metadata`` (the max across DP ranks) so all ranks create
175 + and pick the same buffer at the same step.
176 + """
177 +
178 + def __init__(self, buffer_kwargs: dict[str, Any], max_tokens_per_rank: int):
179 + self._kwargs = {k: v for k, v in buffer_kwargs.items() if k != "S"}
180 + self.max_tokens_per_rank = max_tokens_per_rank
181 + self._buffers: dict[int, Any] = {}
182 +
183 + def capacity_for(self, num_tokens: int) -> int:
184 + if num_tokens >= self.max_tokens_per_rank:
185 + return self.max_tokens_per_rank
186 + cap = 1 << max(num_tokens - 1, 0).bit_length()
187 + return min(max(cap, MOONEP_MIN_DISPATCH_TOKENS), self.max_tokens_per_rank)
188 +
189 + def get(self, num_tokens: int) -> tuple[int, Any]:
190 + """Return ``(capacity, buffer)`` for the given step token count."""
191 + from moonep import Buffer # type: ignore[import-not-found]
192 +
193 + capacity = self.capacity_for(num_tokens)
194 + buffer = self._buffers.get(capacity)
195 + if buffer is None:
196 + buffer = Buffer(S=capacity, **self._kwargs)
197 + self._buffers[capacity] = buffer
198 + return capacity, buffer
199 +
200 + def destroy(self) -> None:
201 + for buffer in self._buffers.values():
202 + buffer.destroy()
203 + self._buffers.clear()
204 +
205 +
206 +class MoonEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
207 + """Prepare/Finalize using MoonEP balanced dispatch/combine.
208 +
209 + ``prepare`` pads the batch to the selected buffer's static token
210 + capacity, dispatches, runs ``prefetch_weight`` for the planned redundant
211 + experts, and stashes the ``plan`` for ``finalize`` (the same pattern
212 + DeepEP-HT uses for its handle). Downstream expert compute must consume
213 + the expert-grouped ``[NvS, H]`` layout via ``cu_seqlens``.
214 + """
215 +
216 + def __init__(
217 + self,
218 + buffer_pool: "MoonEPBufferPool",
219 + max_tokens_per_rank: int,
220 + num_dispatchers: int,
221 + num_global_experts: int,
222 + weight_layout: MoonEPExpertWeightLayout | None = None,
223 + ):
224 + super().__init__()
225 + self.buffer_pool = buffer_pool
226 + self.max_tokens_per_rank = max_tokens_per_rank
227 + self.num_dispatchers_ = num_dispatchers
228 + self.num_global_experts = num_global_experts
229 + self.weight_layout = weight_layout
230 + self._fused_experts: Any = None
231 +
232 + # dispatch state consumed by finalize (and the expert runner)
233 + self._plan: Any = None
234 + self._cu_seqlens: torch.Tensor | None = None
235 + self._num_tokens: int = 0
236 + self._buffer: Any = None # the pool buffer used by the current step
237 +
238 + def post_init_setup(self, fused_experts: mk.FusedMoEExperts) -> None:
239 + # The [E+B] weight layout is attached to the experts by their
240 + # process_weights_after_loading hook (after this runs), so keep the
241 + # reference and resolve the layout lazily in prepare().
242 + self._fused_experts = fused_experts
243 +
244 + def _resolve_weight_layout(self) -> MoonEPExpertWeightLayout:
245 + if self.weight_layout is None and self._fused_experts is not None:
246 + layout = getattr(self._fused_experts, "weight_layout", None)
247 + if isinstance(layout, MoonEPExpertWeightLayout):
248 + self.weight_layout = layout
249 + # Redundant experts' weights must be in rows [E, E+B) before the
250 + # expert compute reads them; skipping prefetch silently corrupts
251 + # output.
252 + assert self.weight_layout is not None, (
253 + "MoonEPPrepareAndFinalize: weight layout not available (the "
254 + "experts' process_weights_after_loading has not run)"
255 + )
256 + return self.weight_layout
257 +
258 + @property
259 + def num_dispatched_slots(self) -> int:
260 + """``NvS`` of the buffer selected by the current step's prepare()."""
261 + assert self._buffer is not None, "prepare() has not been called"
262 + return int(self._buffer._ctx["NvS"])
263 +
264 + def num_dispatchers(self) -> int:
265 + return self.num_dispatchers_
266 +
267 + def output_is_reduced(self) -> bool:
268 + # combine returns the fully weighted+reduced token-major output
269 + return True
270 +
271 + @property
272 + def activation_format(self) -> mk.FusedMoEActivationFormat:
273 + return mk.FusedMoEActivationFormat.Standard
274 +
275 + def max_num_tokens_per_rank(self) -> int | None:
276 + return self.max_tokens_per_rank
277 +
278 + def topk_indices_dtype(self) -> torch.dtype | None:
279 + return torch.int32
280 +
281 + def supports_async(self) -> bool:
282 + return False
283 +
284 + @property
285 + def cu_seqlens(self) -> torch.Tensor:
286 + assert self._cu_seqlens is not None, "prepare() has not been called"
287 + return self._cu_seqlens
288 +
289 + @property
290 + def plan(self) -> Any:
291 + assert self._plan is not None, "prepare() has not been called"
292 + return self._plan
293 +
294 + def _pad_to_capacity(
295 + self,
296 + a1: torch.Tensor,
297 + topk_ids: torch.Tensor,
298 + topk_weights: torch.Tensor,
299 + capacity: int,
300 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
301 + num_tokens = a1.shape[0]
302 + if num_tokens > capacity:
303 + raise ValueError(
304 + f"MoonEP static buffer capacity exceeded: {num_tokens} tokens "
305 + f"> S={capacity}."
306 + )
307 + a1 = a1.contiguous()
308 + topk_ids = topk_ids.to(dtype=torch.int32)
309 + topk_weights = topk_weights.to(dtype=torch.float32)
310 + # MoonEP's planner has no sentinel for invalid expert ids, so map
311 + # negative ids (padded/invalid slots from the engine) to expert 0
312 + # with zero route weight — the same convention as the capacity
313 + # padding rows below. Wasted compute, never wrong output.
314 + invalid = topk_ids < 0
315 + topk_ids = torch.where(invalid, torch.zeros_like(topk_ids), topk_ids)
316 + topk_weights = torch.where(
317 + invalid, torch.zeros_like(topk_weights), topk_weights
318 + ).contiguous()
319 + topk_ids = topk_ids.contiguous()
320 + if num_tokens == capacity:
321 + return a1, topk_ids, topk_weights, num_tokens
322 + pad = capacity - num_tokens
323 + return (
324 + F.pad(a1, (0, 0, 0, pad)),
325 + F.pad(topk_ids, (0, 0, 0, pad)),
326 + F.pad(topk_weights, (0, 0, 0, pad)),
327 + num_tokens,
328 + )
329 +
330 + def prepare(
331 + self,
332 + a1: torch.Tensor,
333 + topk_weights: torch.Tensor,
334 + topk_ids: torch.Tensor,
335 + num_experts: int,
336 + expert_map: torch.Tensor | None,
337 + apply_router_weight_on_input: bool,
338 + quant_config: FusedMoEQuantConfig,
339 + defer_input_quant: bool = False,
340 + ) -> mk.PrepareResultType:
341 + if a1.dtype != torch.bfloat16:
342 + raise NotImplementedError("MoonEP PoC supports BF16 hidden states only.")
343 + if quant_config.quant_dtype is not None:
344 + raise NotImplementedError("MoonEP PoC does not support quantized dispatch.")
345 + assert num_experts == self.num_global_experts
346 + assert self._plan is None, (
347 + "MoonEPPrepareAndFinalize.prepare() called again before finalize()"
348 + )
349 + if apply_router_weight_on_input:
350 + assert topk_ids.size(1) == 1, (
351 + "apply_router_weight_on_input requires top-1 routing"
352 + )
353 + a1 = a1 * topk_weights.to(a1.dtype)
354 +
355 + # Buffer choice must be rank-symmetric: use the max token count
356 + # across DP ranks (every rank sees the same dp_metadata), falling
357 + # back to the local count outside a forward context (tests).
358 + num_tokens = a1.shape[0]
359 + try:
360 + dp_meta = get_forward_context().dp_metadata
361 + except AssertionError:
362 + dp_meta = None
363 + if dp_meta is not None:
364 + num_tokens = int(dp_meta.num_tokens_across_dp_cpu.max())
365 + capacity, self._buffer = self.buffer_pool.get(num_tokens)
366 +
367 + a1, topk_ids, topk_weights, self._num_tokens = self._pad_to_capacity(
368 + a1, topk_ids, topk_weights, capacity
369 + )
370 + tokens_per_expert = torch.bincount(
371 + topk_ids.reshape(-1).to(dtype=torch.int64),
372 + minlength=num_experts,
373 + ).to(dtype=torch.int32)
374 +
375 + hidden_nvsh, route_weights_nvs, cu_seqlens, plan = self._buffer.dispatch(
376 + a1,
377 + topk_weights,
378 + topk_ids,
379 + tokens_per_expert,
380 + )
381 + self._plan = plan
382 + self._cu_seqlens = cu_seqlens
383 +
384 + weight_layout = self._resolve_weight_layout()
385 + self._buffer.prefetch_weight(
386 + plan=plan,
387 + full_gate_weight=weight_layout.full_gate_weight,
388 + full_up_weight=weight_layout.full_up_weight,
389 + full_down_weight=weight_layout.full_down_weight,
390 + )
391 +
392 + # Segment sizes per [E+B] weight row. NOTE: this is per *row*, not
393 + # per local expert, and NvS rows are expert-grouped rather than
394 + # token-major — only MoonEP-aware expert implementations can consume
395 + # this activation layout.
396 + expert_num_tokens = torch.diff(cu_seqlens, prepend=cu_seqlens.new_zeros(1))
397 + expert_tokens_meta = mk.ExpertTokensMetadata(
398 + expert_num_tokens=expert_num_tokens,
399 + expert_num_tokens_cpu=None,
400 + )
401 +
402 + # MoonEP has no post-dispatch per-token topk ids/weights; route
403 + # weights come back in NvS order.
404 + return hidden_nvsh, None, expert_tokens_meta, None, route_weights_nvs
405 +
406 + def finalize(
407 + self,
408 + output: torch.Tensor,
409 + fused_expert_output: torch.Tensor,
410 + topk_weights: torch.Tensor,
411 + topk_ids: torch.Tensor,
412 + apply_router_weight_on_input: bool,
413 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
414 + ) -> None:
415 + # Route weights are applied in the expert compute and MoonEP's
416 + # combine performs the K-sum, so any weight/reduce work in finalize
417 + # means MoonEP was paired with the wrong kind of experts.
418 + assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP), (
419 + "MoonEP requires TopKWeightAndReduceNoOP, got "
420 + f"{type(weight_and_reduce_impl).__name__}"
421 + )
422 + combined, _, _ = self._buffer.combine(
423 + plan=self.plan,
424 + hidden_nvsh=fused_expert_output,
425 + route_weights_nvs=None,
426 + )
427 + output.copy_(combined[: self._num_tokens])
428 + self._plan = None
429 + self._cu_seqlens = None
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+14 -10
@@ -49,13 +49,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
49
50 # --8<-- [end:unquantized_fused_moe]
51
52 - supports_pre_processed_weights = True
53 -
52 def __init__(self, moe: FusedMoEConfig):
53 super().__init__(moe)
54 self.unquantized_backend, self.experts_cls = select_unquantized_moe_backend(
55 moe_config=self.moe,
56 )
57 + # MoonEP replaces the named parameters with its [E+B] gate/down
58 + # split and keeps the up projection outside the parameter registry,
59 + # so a pre-processed (e.g. ipc_cache) load cannot reconstruct it.
60 + self.supports_pre_processed_weights = (
61 + self.unquantized_backend != UnquantizedMoeBackend.MOONEP
62 + )
63
64 @property
65 def supports_eplb(self) -> bool:
@@ -171,6 +175,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
175 moe_config=layer.moe_config,
176 w13_weight=w13,
177 w2_weight=w2,
178 + layer=layer,
179 )
180 # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__;
181 # On the first call we replace the parameter normally. On subsequent
@@ -191,14 +196,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
196 # does not need to be re-built.
197 self._init_moe_kernel(layer)
198
194 - if self.unquantized_backend == UnquantizedMoeBackend.CPU:
195 - # The CPU experts need the layer itself for the setup that
196 - # convert_to_unquantized_kernel_format cannot express, since
197 - # it only sees the two weight tensors: padding and prepacking
198 - # into the grouped-gemm layout (bias included), and capturing
199 - # the router config that monolithic apply() cannot carry.
200 - assert self.moe_kernel is not None
201 - self.moe_kernel.fused_experts.process_weights_after_loading(layer)
199 + # No-op by default. Experts that need the layer itself for setup
200 + # convert_to_unquantized_kernel_format cannot express override
201 + # it: CPU prepacks into its grouped-gemm layout and captures the
202 + # router config; MoonEP picks up the [E+B] weight layout for
203 + # prefetch and the up projection.
204 + assert self.moe_kernel is not None
205 + self.moe_kernel.fused_experts.process_weights_after_loading(layer)
206
207 def _init_moe_kernel(self, layer: "RoutedExperts") -> None:
208 """Build the MoE kernel from the layer's current (shuffled) weights."""
vllm/utils/import_utils.py
+28
@@ -512,6 +512,34 @@ def has_nixl_ep() -> bool:
512 return _has_module("nixl_ep")
513
514
515 +def has_moonep() -> bool:
516 + """Whether the optional `moonep` package is available."""
517 + return _has_module("moonep")
518 +
519 +
520 +def check_moonep_system_support() -> None:
521 + """Raise if the current device cannot run MoonEP.
522 +
523 + MoonEP's symmetric-memory buffers require CUDA VMM plus NVSwitch multicast
524 + (SHARP) on every EP rank; NVLink-only topologies without NVSwitch (e.g.
525 + 4x H100 NV6) fail deep inside ``moonep.Buffer`` otherwise.
526 + """
527 + if not has_moonep():
528 + raise RuntimeError(
529 + "MoonEP not available. Install it from "
530 + "https://github.com/MoonshotAI/MoonEP."
531 + )
532 + from moonep._C import ( # type: ignore[import-not-found]
533 + nvl_multicast_supported,
534 + )
535 +
536 + if not nvl_multicast_supported():
537 + raise RuntimeError(
538 + "MoonEP requires NVSwitch multicast (SHARP) support on the "
539 + "current device; it is not supported on this GPU/topology."
540 + )
541 +
542 +
543 def is_numba_available() -> bool:
544 """Whether the optional `numba` package is available."""
545 return _has_module("numba")