| 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 |