| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the MOE layer. |
| 4 | |
| 5 | Run `pytest tests/kernels/test_moe_layer.py`. |
| 6 | """ |
| 7 | |
| 8 | import functools |
| 9 | import os |
| 10 | import sys |
| 11 | import tempfile |
| 12 | import traceback |
| 13 | import types |
| 14 | from collections.abc import Callable |
| 15 | from contextlib import suppress |
| 16 | from dataclasses import astuple, dataclass, fields, replace |
| 17 | from itertools import product |
| 18 | from typing import get_args |
| 19 | |
| 20 | import pytest |
| 21 | import torch |
| 22 | |
| 23 | import vllm.model_executor.layers.quantization.utils.w8a8_utils |
| 24 | from tests.kernels.moe.modular_kernel_tools.parallel_utils import ( |
| 25 | ProcessGroupInfo, |
| 26 | _set_vllm_config, |
| 27 | parallel_launch_with_config, |
| 28 | ) |
| 29 | from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_weights |
| 30 | from vllm.config import ( |
| 31 | CompilationConfig, |
| 32 | EPLBConfig, |
| 33 | ParallelConfig, |
| 34 | SchedulerConfig, |
| 35 | VllmConfig, |
| 36 | set_current_vllm_config, |
| 37 | ) |
| 38 | from vllm.distributed import ( |
| 39 | get_ep_group, |
| 40 | get_eplb_group, |
| 41 | tensor_model_parallel_all_gather, |
| 42 | ) |
| 43 | from vllm.distributed.device_communicators.all_reduce_utils import ( |
| 44 | gpu_p2p_access_check, |
| 45 | ) |
| 46 | from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator |
| 47 | from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace |
| 48 | from vllm.forward_context import set_forward_context |
| 49 | from vllm.model_executor.layers.fused_moe import ( |
| 50 | FusedMoEFactory, |
| 51 | MoERunner, |
| 52 | fused_experts, |
| 53 | ) |
| 54 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 55 | from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig |
| 56 | from vllm.model_executor.layers.fused_moe.router.router_factory import ( |
| 57 | create_fused_moe_router, |
| 58 | ) |
| 59 | from vllm.model_executor.layers.quantization.base_config import QuantizationConfig |
| 60 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 61 | ModelOptFp8Config, |
| 62 | ModelOptNvFp4Config, |
| 63 | ) |
| 64 | from vllm.model_executor.models.utils import sequence_parallel_chunk |
| 65 | from vllm.platforms import current_platform |
| 66 | from vllm.utils.flashinfer import ( |
| 67 | has_flashinfer_nvlink_one_sided, |
| 68 | has_flashinfer_nvlink_two_sided, |
| 69 | ) |
| 70 | from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep |
| 71 | from vllm.utils.math_utils import cdiv, next_power_of_2 |
| 72 | from vllm.utils.torch_utils import set_random_seed |
| 73 | from vllm.v1.worker.workspace import ( |
| 74 | init_workspace_manager, |
| 75 | is_workspace_manager_initialized, |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | def on_gfx950() -> bool: |
| 80 | if current_platform.is_rocm(): |
| 81 | from vllm.platforms.rocm import on_gfx950 as rocm_on_gfx950 |
| 82 | |
| 83 | return rocm_on_gfx950() |
| 84 | return False |
| 85 | |
| 86 | |
| 87 | fp8_dtype = current_platform.fp8_dtype() |
| 88 | |
| 89 | SHAPE_COMBOS = [ |
| 90 | (1, 128, 256), |
| 91 | (32, 512, 512), |
| 92 | (222, 1024, 2048), |
| 93 | ] |
| 94 | MAX_M = max([x[0] for x in SHAPE_COMBOS]) |
| 95 | |
| 96 | NUM_EXPERTS = [8, 64] |
| 97 | TOP_KS = [2, 6] |
| 98 | |
| 99 | # dp_size, tp_size, use_ep |
| 100 | # Note: DP+TP is not yet supported in the FusedMoEFactory layer. |
| 101 | PARALLEL_COMBOS = [ |
| 102 | [1, 2, False], |
| 103 | [1, 4, False], |
| 104 | [2, 1, True], |
| 105 | [4, 1, True], |
| 106 | # This combination indicates sequence parallel. |
| 107 | # See ParallelConfig.use_sequence_parallel. |
| 108 | [2, 2, True], |
| 109 | ] |
| 110 | |
| 111 | # TODO: should this even be set manually? let oracles handle this |
| 112 | BACKENDS = ["allgather_reducescatter"] |
| 113 | |
| 114 | if has_mori(): |
| 115 | BACKENDS += ["mori_high_throughput", "mori_low_latency"] |
| 116 | |
| 117 | if has_flashinfer_nvlink_two_sided(): |
| 118 | BACKENDS += ["flashinfer_nvlink_two_sided"] |
| 119 | |
| 120 | if has_flashinfer_nvlink_one_sided(): |
| 121 | BACKENDS += ["flashinfer_nvlink_one_sided"] |
| 122 | |
| 123 | if has_deep_ep(): |
| 124 | BACKENDS += ["deepep_high_throughput", "deepep_low_latency"] |
| 125 | |
| 126 | if has_nixl_ep(): |
| 127 | BACKENDS += ["nixl_ep"] |
| 128 | |
| 129 | DEEPEP_BACKENDS = {"deepep_high_throughput", "deepep_low_latency"} |
| 130 | MORI_BACKENDS = {"mori_high_throughput", "mori_low_latency"} |
| 131 | |
| 132 | QUANT_METHODS = [ |
| 133 | None, |
| 134 | "fp8", |
| 135 | "fp8_blocked", |
| 136 | "modelopt_fp8", |
| 137 | "modelopt_fp4", |
| 138 | ] |
| 139 | |
| 140 | # Which quantization methods each backend supports. |
| 141 | # fmt: off |
| 142 | BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = { |
| 143 | "allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501 |
| 144 | "mori_high_throughput": {None, "fp8", "modelopt_fp8"}, |
| 145 | "mori_low_latency": {None, "fp8", "modelopt_fp8"}, |
| 146 | "flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 |
| 147 | "flashinfer_nvlink_one_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 |
| 148 | "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 |
| 149 | "deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501 |
| 150 | "nixl_ep": {None, "fp8_blocked", "modelopt_fp8"}, |
| 151 | } |
| 152 | |
| 153 | # Map from backend -> (DP/EP support, DP support, TP support, SP support) |
| 154 | BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool, bool]] = { |
| 155 | "allgather_reducescatter": (True, True, True, True), |
| 156 | "mori_high_throughput": (True, False, False, True), |
| 157 | "mori_low_latency": (True, False, False, True), |
| 158 | "flashinfer_nvlink_two_sided": (False, True, False, False), |
| 159 | "flashinfer_nvlink_one_sided": (False, True, False, False), |
| 160 | "deepep_low_latency": (True, False, False, True), |
| 161 | "deepep_high_throughput": (True, False, False, True), |
| 162 | "nixl_ep": (True, False, False, True), |
| 163 | } |
| 164 | # fmt: on |
| 165 | |
| 166 | # Which quantization methods support EPLB. |
| 167 | # ModelOptFp8MoEMethod inherits supports_eplb=False from FusedMoEMethodBase. |
| 168 | # TODO: double check modelopt fp8 |
| 169 | # modelopt_fp4 excluded: get_expert_weights() can't handle NvFP4 packed format. |
| 170 | EPLB_SUPPORTED_QUANTS: list[str | None] = [None, "fp8"] |
| 171 | |
| 172 | # Which backends support EPLB. |
| 173 | # deepep backends fail in get_expert_weights / rearrange_expert_weights_inplace. |
| 174 | # TODO(bnell): check this |
| 175 | EPLB_SUPPORTED_BACKENDS: list[str] = ["allgather_reducescatter"] |
| 176 | |
| 177 | |
| 178 | def mock_normalize_e4m3fn_to_e4m3fnuz( |
| 179 | weight: torch.Tensor, |
| 180 | weight_scale: torch.Tensor, |
| 181 | input_scale: torch.Tensor | None = None, |
| 182 | ): |
| 183 | return weight, weight_scale, input_scale |
| 184 | |
| 185 | |
| 186 | # Needed since weights will already be in e4m3fnuz format on platforms that |
| 187 | # use the fnuz fp8 format and the normalize_e4m3fn_to_e4m3fnuz() function |
| 188 | # is not being tested here. |
| 189 | # NOTE: The weights are quantized by moe_quantize_weights_2d in |
| 190 | # _quantize_fp8_halves. |
| 191 | # NOTE: Not able to use monkeypatch because of the spawned parallel workers. |
| 192 | def override_normalize_e4m3fn_to_e4m3fnuz(): |
| 193 | vllm.model_executor.layers.quantization.utils.w8a8_utils.normalize_e4m3fn_to_e4m3fnuz = mock_normalize_e4m3fn_to_e4m3fnuz # noqa: E501 |
| 194 | vllm.model_executor.layers.quantization.fp8.normalize_e4m3fn_to_e4m3fnuz = ( |
| 195 | mock_normalize_e4m3fn_to_e4m3fnuz # noqa: E501 |
| 196 | ) |
| 197 | |
| 198 | |
| 199 | def sp_wrapper( |
| 200 | fn: Callable | MoERunner, is_sequence_parallel: bool | None = None |
| 201 | ) -> Callable: |
| 202 | """Wrapper to handle sequence parallelism chunking and gathering. |
| 203 | |
| 204 | For SP with EP: |
| 205 | - The TP group is created with the original tensor_parallel_size (e.g., 2) |
| 206 | - get_tp_group() has the correct world_size for SP operations |
| 207 | - sequence_parallel_chunk() uses get_tensor_model_parallel_world_size() |
| 208 | - tensor_model_parallel_all_gather() uses get_tp_group() |
| 209 | - Both should work correctly even when EP is enabled |
| 210 | """ |
| 211 | if isinstance(fn, MoERunner): |
| 212 | assert is_sequence_parallel is None |
| 213 | is_sequence_parallel = fn.moe_config.moe_parallel_config.is_sequence_parallel |
| 214 | else: |
| 215 | assert is_sequence_parallel is not None |
| 216 | |
| 217 | if is_sequence_parallel: |
| 218 | |
| 219 | def wrapper( |
| 220 | hidden_states: torch.Tensor, |
| 221 | router_logits: torch.Tensor, |
| 222 | ) -> torch.Tensor: |
| 223 | # Split sequence across TP ranks |
| 224 | # Both hidden_states and router_logits have [num_tokens, ...] shape |
| 225 | hidden_states = sequence_parallel_chunk(hidden_states) |
| 226 | router_logits = sequence_parallel_chunk(router_logits) |
| 227 | # Run MoE on local chunk |
| 228 | result = fn(hidden_states, router_logits) |
| 229 | # Gather results from all TP ranks |
| 230 | result = tensor_model_parallel_all_gather(result, 0) |
| 231 | # Remove any padding added by SP. |
| 232 | return result[: hidden_states.shape[0]] |
| 233 | |
| 234 | return wrapper |
| 235 | return fn |
| 236 | |
| 237 | |
| 238 | def maybe_roundup_layer_hidden_size( |
| 239 | hidden_size: int, |
| 240 | act_dtype: torch.dtype, |
| 241 | backend: str | None, |
| 242 | ) -> int: |
| 243 | """Given layer hidden size and MoE configurations, round up hidden_size |
| 244 | if necessary. |
| 245 | |
| 246 | Args: |
| 247 | hidden_size: Layer hidden-size |
| 248 | act_dtype: Data type of the layer activations. |
| 249 | moe_parallel_config: Fused MoE parallelization strategy configuration. |
| 250 | |
| 251 | Return: |
| 252 | Rounded up hidden_size if rounding up is required based on the configs |
| 253 | and all2all backend. |
| 254 | Original hidden size otherwise. |
| 255 | |
| 256 | """ |
| 257 | if backend == "deepep_high_throughput": |
| 258 | from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import ( |
| 259 | DeepEPHTPrepareAndFinalize, |
| 260 | ) |
| 261 | |
| 262 | hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size( |
| 263 | hidden_size, act_dtype |
| 264 | ) |
| 265 | |
| 266 | if backend == "deepep_low_latency": |
| 267 | from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( |
| 268 | DeepEPLLPrepareAndFinalize, |
| 269 | ) |
| 270 | |
| 271 | hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size( |
| 272 | hidden_size |
| 273 | ) |
| 274 | |
| 275 | return hidden_size |
| 276 | |
| 277 | |
| 278 | def chunk_by_rank( |
| 279 | t: torch.Tensor, |
| 280 | r: int, |
| 281 | w: int, |
| 282 | dim: int = 0, |
| 283 | device: torch.device | None = None, |
| 284 | ) -> torch.Tensor: |
| 285 | chunk = cdiv(t.shape[dim], w) |
| 286 | t = t.narrow(dim, r * chunk, chunk) |
| 287 | if device is not None: |
| 288 | t = t.to(device) |
| 289 | return t |
| 290 | |
| 291 | |
| 292 | def tp_chunk_gate_up( |
| 293 | w: torch.Tensor, |
| 294 | tp_rank: int, |
| 295 | tp_size: int, |
| 296 | dim: int, |
| 297 | device: torch.device | int | None = None, |
| 298 | ) -> torch.Tensor: |
| 299 | """TP-chunk a combined [gate; up] weight, splitting each half separately |
| 300 | so every rank gets a portion of both gate and up.""" |
| 301 | half = w.shape[dim] // 2 |
| 302 | gate = chunk_by_rank( |
| 303 | w.narrow(dim, 0, half), tp_rank, tp_size, dim=dim, device=device |
| 304 | ) |
| 305 | up = chunk_by_rank( |
| 306 | w.narrow(dim, half, half), tp_rank, tp_size, dim=dim, device=device |
| 307 | ) |
| 308 | return torch.cat([gate, up], dim=dim) |
| 309 | |
| 310 | |
| 311 | @dataclass |
| 312 | class MoETestConfig: |
| 313 | m: int |
| 314 | n: int |
| 315 | k: int |
| 316 | num_experts: int |
| 317 | top_k: int |
| 318 | in_dtype: torch.dtype |
| 319 | quantization: str | None |
| 320 | use_shared_experts: bool |
| 321 | use_gate: bool |
| 322 | use_routed_input_transform: bool |
| 323 | enable_eplb: bool = False |
| 324 | backend: str | None = None |
| 325 | ep_size: int = 1 |
| 326 | dp_size: int = 1 |
| 327 | tp_size: int = 1 |
| 328 | |
| 329 | @property |
| 330 | def is_sequence_parallel(self) -> bool: |
| 331 | # Sequence parallelism: EP enabled + TP dimension used for sequence splitting |
| 332 | # In test config: ep_size represents total expert parallel size |
| 333 | # tp_size represents the original TP dimension (becomes sp_size in MoERunner) |
| 334 | # dp_size represents data parallel size |
| 335 | # For SP: we need EP enabled (ep_size > 1) and sequence splitting (tp_size > 1) |
| 336 | return self.ep_size > 1 and self.tp_size > 1 |
| 337 | |
| 338 | # TODO: add more error messages |
| 339 | def id(self) -> str: |
| 340 | def proc(s: str) -> str: |
| 341 | return s.removeprefix("torch.") |
| 342 | |
| 343 | id_str = "-".join([proc(str(f)) for f in astuple(self)]) |
| 344 | return f"[{id_str}]" |
| 345 | |
| 346 | # TODO: add more error messages |
| 347 | @staticmethod |
| 348 | def from_id(id: str) -> "MoETestConfig": |
| 349 | id = id[1:-1] |
| 350 | str_values = id.split("-") |
| 351 | |
| 352 | def convert(v: str, ty): |
| 353 | if isinstance(ty, types.UnionType): |
| 354 | sub_ty = list(get_args(ty)) |
| 355 | assert len(sub_ty) == 2 and types.NoneType in sub_ty |
| 356 | sub_ty.remove(types.NoneType) |
| 357 | return sub_ty[0](v) if v != "None" else None |
| 358 | elif ty is torch.dtype: |
| 359 | ty_val = getattr(torch, v, None) |
| 360 | assert isinstance(ty_val, torch.dtype) |
| 361 | return ty_val |
| 362 | elif ty is bool: |
| 363 | return v == "True" |
| 364 | else: |
| 365 | return ty(v) |
| 366 | |
| 367 | values = tuple( |
| 368 | [convert(v, f.type) for v, f in zip(str_values, fields(MoETestConfig))] |
| 369 | ) |
| 370 | return MoETestConfig(*values) |
| 371 | |
| 372 | |
| 373 | def _group_deepep_ll_configs( |
| 374 | test_configs: list[MoETestConfig], |
| 375 | ) -> list[list[MoETestConfig]]: |
| 376 | """Group configs that can share one DeepEP low-latency buffer.""" |
| 377 | groups: dict[tuple[int, int, int], list[MoETestConfig]] = {} |
| 378 | for test_config in test_configs: |
| 379 | # max_num_tokens and the remaining buffer arguments are fixed for one |
| 380 | # outer pytest item. These are the arguments that can vary here. |
| 381 | buffer_key = ( |
| 382 | test_config.k, |
| 383 | test_config.ep_size, |
| 384 | test_config.num_experts, |
| 385 | ) |
| 386 | groups.setdefault(buffer_key, []).append(test_config) |
| 387 | return list(groups.values()) |
| 388 | |
| 389 | |
| 390 | @pytest.mark.cpu_test |
| 391 | def test_group_deepep_ll_configs_by_buffer_requirements(): |
| 392 | config_a = MoETestConfig( |
| 393 | 1, |
| 394 | 128, |
| 395 | 2048, |
| 396 | 8, |
| 397 | 2, |
| 398 | torch.bfloat16, |
| 399 | None, |
| 400 | False, |
| 401 | False, |
| 402 | False, |
| 403 | backend="deepep_low_latency", |
| 404 | ep_size=2, |
| 405 | dp_size=2, |
| 406 | ) |
| 407 | config_b = replace(config_a, num_experts=64) |
| 408 | config_a_later = replace(config_a, m=32) |
| 409 | config_c = replace(config_a, k=4096) |
| 410 | config_d = replace(config_a, ep_size=4, dp_size=4) |
| 411 | |
| 412 | assert _group_deepep_ll_configs( |
| 413 | [config_a, config_b, config_a_later, config_c, config_d] |
| 414 | ) == [ |
| 415 | [config_a, config_a_later], |
| 416 | [config_b], |
| 417 | [config_c], |
| 418 | [config_d], |
| 419 | ] |
| 420 | assert _group_deepep_ll_configs([]) == [] |
| 421 | |
| 422 | |
| 423 | def generate_valid_test_configs( |
| 424 | backend: str, |
| 425 | ep_size: int, |
| 426 | dp_size: int, |
| 427 | tp_size: int, |
| 428 | enable_eplb: bool, |
| 429 | verbosity: int = 0, |
| 430 | ) -> list[MoETestConfig]: |
| 431 | configs: list[MoETestConfig] = [] |
| 432 | |
| 433 | for ( |
| 434 | shape, |
| 435 | num_experts, |
| 436 | top_k, |
| 437 | quantization, |
| 438 | use_shared_experts, |
| 439 | use_gate, |
| 440 | use_routed_input_transform, |
| 441 | ) in product( |
| 442 | SHAPE_COMBOS, |
| 443 | NUM_EXPERTS, |
| 444 | TOP_KS, |
| 445 | QUANT_METHODS, |
| 446 | [False, True], # shared |
| 447 | [False, True], # gate |
| 448 | [False, True], # routed input exform |
| 449 | ): |
| 450 | config = MoETestConfig( |
| 451 | shape[0], # m |
| 452 | shape[1], # n |
| 453 | shape[2], # k |
| 454 | num_experts, |
| 455 | top_k, |
| 456 | torch.bfloat16, |
| 457 | quantization, |
| 458 | use_shared_experts, |
| 459 | use_gate, |
| 460 | use_routed_input_transform, |
| 461 | enable_eplb, |
| 462 | backend, |
| 463 | ep_size, |
| 464 | dp_size, |
| 465 | tp_size, |
| 466 | ) |
| 467 | |
| 468 | valid, reason = is_valid_config(config) |
| 469 | if valid: |
| 470 | configs.append(config) |
| 471 | elif verbosity > 1: |
| 472 | print(f"Skipping invalid config {config} - {reason}") |
| 473 | |
| 474 | return configs |
| 475 | |
| 476 | |
| 477 | @functools.cache |
| 478 | def visible_devices_have_peer_access(world_size: int) -> bool: |
| 479 | if not current_platform.is_cuda(): |
| 480 | return True |
| 481 | |
| 482 | try: |
| 483 | return all( |
| 484 | gpu_p2p_access_check(src, dst) |
| 485 | for src in range(world_size) |
| 486 | for dst in range(world_size) |
| 487 | if src != dst |
| 488 | ) |
| 489 | except RuntimeError: |
| 490 | return False |
| 491 | |
| 492 | |
| 493 | # TODO: break this up into sections |
| 494 | def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: |
| 495 | # routed_input_transform only makes sense with shared_experts (latent MoE) |
| 496 | # TODO: not sure this is true |
| 497 | if config.use_routed_input_transform and not config.use_shared_experts: |
| 498 | return False, "routed_input_transform requires shared_experts" |
| 499 | |
| 500 | # TODO: disable for now |
| 501 | if config.use_routed_input_transform and config.enable_eplb: |
| 502 | return False, "routed_input_transform not supported with EPLB." |
| 503 | |
| 504 | # TODO: disable for now |
| 505 | if config.use_routed_input_transform and config.use_gate: |
| 506 | return ( |
| 507 | False, |
| 508 | "routed_input_transform not supported with gate because of " |
| 509 | "padding problems", |
| 510 | ) |
| 511 | |
| 512 | # TODO: disable for now |
| 513 | if config.use_routed_input_transform and config.backend in [ |
| 514 | "deepep_low_latency", |
| 515 | "deepep_high_throughput", |
| 516 | ]: |
| 517 | return ( |
| 518 | False, |
| 519 | "routed_input_transform not supported with DeepEP backends because " |
| 520 | "of padding problems", |
| 521 | ) |
| 522 | |
| 523 | # routed_input_transform + quantization + high hidden dimensions |
| 524 | # TODO: Disable >= 2048 for now due to insane errors. |
| 525 | if ( |
| 526 | config.use_routed_input_transform |
| 527 | and config.quantization is not None |
| 528 | and config.k >= 2048 |
| 529 | ): |
| 530 | return ( |
| 531 | False, |
| 532 | "routed_input_transform + quantization + higher hidden dimensions " |
| 533 | "leads to large differences.", |
| 534 | ) |
| 535 | |
| 536 | # Skip modelopt_fp4 if not on B100+ (compute capability 10.0+) or gfx950. |
| 537 | if config.quantization == "modelopt_fp4" and not ( |
| 538 | (current_platform.is_rocm() and on_gfx950()) |
| 539 | or current_platform.has_device_capability(100) |
| 540 | ): |
| 541 | return False, "modelopt_fp4 requires native NVFP4 or emulation" |
| 542 | |
| 543 | # Skip flashinfer_nvlink if not on H100+ (compute capability 10.0+) |
| 544 | if ( |
| 545 | config.backend is not None |
| 546 | and config.backend.startswith("flashinfer_nvlink") |
| 547 | and not current_platform.has_device_capability(90) |
| 548 | ): |
| 549 | return False, "flashinfer_nvlink needs H100+ GPUs" |
| 550 | |
| 551 | # Backend-specific checks |
| 552 | if config.backend is not None: |
| 553 | supported_quants = BACKEND_SUPPORTED_QUANTS.get(config.backend) |
| 554 | if supported_quants is not None and config.quantization not in supported_quants: |
| 555 | return ( |
| 556 | False, |
| 557 | f"{config.backend} does not support quantization={config.quantization}", |
| 558 | ) |
| 559 | |
| 560 | if ( |
| 561 | on_gfx950() |
| 562 | and config.backend == "deepep_low_latency" |
| 563 | and config.quantization == "modelopt_fp4" |
| 564 | ): |
| 565 | return ( |
| 566 | False, |
| 567 | "DeepEP low latency requires a batched NVFP4 MoE backend, " |
| 568 | "which is not available on gfx950.", |
| 569 | ) |
| 570 | |
| 571 | if config.backend in MORI_BACKENDS: |
| 572 | if os.environ.get("VLLM_TEST_ENABLE_MORI_MOE_LAYER") != "1": |
| 573 | return False, "mori MoE layer matrix is opt-in" |
| 574 | |
| 575 | from vllm._aiter_ops import rocm_aiter_ops |
| 576 | |
| 577 | if not rocm_aiter_ops.is_fused_moe_enabled(): |
| 578 | return False, "mori requires AITER fused MoE" |
| 579 | if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled(): |
| 580 | return False, "mori does not support AITER shared expert fusion" |
| 581 | |
| 582 | if config.backend == "deepep_low_latency": |
| 583 | from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( # noqa: E501 |
| 584 | DeepEPLLPrepareAndFinalize, |
| 585 | ) |
| 586 | |
| 587 | if config.k not in DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES: |
| 588 | return ( |
| 589 | False, |
| 590 | f"Skipping unsupported K {config.k} in {config.backend} w/o EP.", |
| 591 | ) |
| 592 | |
| 593 | if config.backend == "nixl_ep": |
| 594 | from vllm.model_executor.layers.fused_moe.prepare_finalize.nixl_ep import ( # noqa: E501 |
| 595 | NixlEPPrepareAndFinalize, |
| 596 | ) |
| 597 | |
| 598 | if config.k not in NixlEPPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES: |
| 599 | return ( |
| 600 | False, |
| 601 | f"Skipping unsupported K {config.k} in {config.backend} w/o EP.", |
| 602 | ) |
| 603 | |
| 604 | if config.backend is not None: |
| 605 | supports_ep_dp, supports_dp, supports_tp, supports_sp = ( |
| 606 | BACKEND_EP_DP_TP_SUPPORT[config.backend] |
| 607 | ) |
| 608 | |
| 609 | if config.tp_size > 1 and not supports_tp and not config.is_sequence_parallel: |
| 610 | return False, f"{config.backend} does not support TP." |
| 611 | |
| 612 | if config.dp_size > 1 and config.ep_size == 1 and not supports_dp: |
| 613 | return False, f"{config.backend} does not support DP." |
| 614 | |
| 615 | if config.dp_size > 1 and config.ep_size > 1 and not supports_ep_dp: |
| 616 | return False, f"{config.backend} does not support EP/DP." |
| 617 | |
| 618 | if config.is_sequence_parallel and not supports_sp: |
| 619 | return False, f"{config.backend} does not support SP." |
| 620 | else: |
| 621 | if config.tp_size > 1 or config.ep_size > 1 or config.dp_size > 1: |
| 622 | return False, "An all2all backend is required for parallelism." |
| 623 | |
| 624 | # Sequence parallelism specific validations |
| 625 | if config.is_sequence_parallel: |
| 626 | if config.ep_size == 1: |
| 627 | return False, "Sequence parallelism requires EP to be enabled (ep_size > 1)" |
| 628 | |
| 629 | if config.tp_size == 1: |
| 630 | return ( |
| 631 | False, |
| 632 | "Sequence parallelism requires tp_size > 1 for sequence splitting", |
| 633 | ) |
| 634 | |
| 635 | # SP is essentially EP + sequence splitting |
| 636 | # Verify the relationship: ep_size should equal dp_size * tp_size |
| 637 | # (when pcp_size=1). |
| 638 | expected_ep_size = config.dp_size * config.tp_size |
| 639 | if config.ep_size != expected_ep_size: |
| 640 | return False, ( |
| 641 | f"For sequence parallelism: ep_size ({config.ep_size}) should equal " |
| 642 | f"dp_size * tp_size ({expected_ep_size})" |
| 643 | ) |
| 644 | |
| 645 | if config.enable_eplb: |
| 646 | if config.ep_size == 1: |
| 647 | return False, "EPLB requires EP." |
| 648 | |
| 649 | if config.quantization not in EPLB_SUPPORTED_QUANTS: |
| 650 | return False, f"EPLB not supported with {config.quantization} quantization." |
| 651 | |
| 652 | if config.backend not in EPLB_SUPPORTED_BACKENDS: |
| 653 | return False, f"EPLB not supported with {config.backend}." |
| 654 | |
| 655 | if config.num_experts % config.dp_size != 0: |
| 656 | return False, "EPLB requires num_experts divisible by ep_size" |
| 657 | |
| 658 | return True, None |
| 659 | |
| 660 | |
| 661 | def chunk_scales( |
| 662 | t: torch.Tensor | None, |
| 663 | start: int, |
| 664 | end: int, |
| 665 | device: torch.device | None = None, |
| 666 | ) -> torch.Tensor | None: |
| 667 | if t is not None and t.numel() > 1: |
| 668 | t = t[start:end] |
| 669 | |
| 670 | if t is not None and device is not None: |
| 671 | t = t.to(device) |
| 672 | |
| 673 | return t |
| 674 | |
| 675 | |
| 676 | @dataclass |
| 677 | class QuantizedWeights: |
| 678 | w13_weight: torch.Tensor |
| 679 | w2_weight: torch.Tensor |
| 680 | w13_weight_scale: torch.Tensor | None = None |
| 681 | w2_weight_scale: torch.Tensor | None = None |
| 682 | w13_weight_scale_2: torch.Tensor | None = None |
| 683 | w2_weight_scale_2: torch.Tensor | None = None |
| 684 | w13_input_scale: torch.Tensor | None = None |
| 685 | w2_input_scale: torch.Tensor | None = None |
| 686 | |
| 687 | |
| 688 | def _quantize_fp8_halves( |
| 689 | w1: torch.Tensor, |
| 690 | w2: torch.Tensor, |
| 691 | block_shape: list[int] | None = None, |
| 692 | ) -> QuantizedWeights: |
| 693 | """Quantize w13 gate/up halves separately to FP8, producing per-shard scales.""" |
| 694 | half = w1.shape[1] // 2 |
| 695 | w1q_a, w1s_a, _ = moe_quantize_weights( |
| 696 | w1[:, :half, :], |
| 697 | None, |
| 698 | fp8_dtype, |
| 699 | False, |
| 700 | block_shape, |
| 701 | ) |
| 702 | w1q_b, w1s_b, _ = moe_quantize_weights( |
| 703 | w1[:, half:, :], |
| 704 | None, |
| 705 | fp8_dtype, |
| 706 | False, |
| 707 | block_shape, |
| 708 | ) |
| 709 | assert w1s_a is not None and w1s_b is not None |
| 710 | |
| 711 | w2q, w2s, _ = moe_quantize_weights(w2, None, fp8_dtype, False, block_shape) |
| 712 | assert w2s is not None |
| 713 | |
| 714 | if block_shape is not None: |
| 715 | # Blocked quantization: scales have shape (E, n_tiles, k_tiles) |
| 716 | # Concatenate gate and up scales along the n_tiles dimension (dim=1) |
| 717 | # to match the concatenation of gate and up weights |
| 718 | w13_weight_scale = torch.cat([w1s_a, w1s_b], dim=1) |
| 719 | # w2 scales keep their blocked shape (E, k_tiles, n_tiles) |
| 720 | w2_weight_scale = w2s |
| 721 | else: |
| 722 | # Non-blocked quantization: scales have shape (E, 1, 1) |
| 723 | # Each w1s_x is (E, 1, 1) -> reshape to (E, 1), cat to (E, 2) |
| 724 | w13_weight_scale = torch.cat([w1s_a.view(-1, 1), w1s_b.view(-1, 1)], dim=1) |
| 725 | # w2s is (E, 1, 1) -> reshape to (E,) |
| 726 | w2_weight_scale = w2s.view(-1) |
| 727 | |
| 728 | return QuantizedWeights( |
| 729 | w13_weight=torch.cat([w1q_a, w1q_b], dim=1), |
| 730 | w2_weight=w2q, |
| 731 | w13_weight_scale=w13_weight_scale, |
| 732 | w2_weight_scale=w2_weight_scale, |
| 733 | ) |
| 734 | |
| 735 | |
| 736 | def make_quant_config( |
| 737 | quantization: str | None, |
| 738 | w1: torch.Tensor, |
| 739 | w2: torch.Tensor, |
| 740 | num_experts: int, |
| 741 | ) -> tuple[QuantizationConfig | None, QuantizedWeights]: |
| 742 | from vllm.model_executor.layers.quantization.fp8 import Fp8Config |
| 743 | |
| 744 | if quantization is None: |
| 745 | return None, QuantizedWeights(w13_weight=w1, w2_weight=w2) |
| 746 | |
| 747 | if quantization == "fp8": |
| 748 | return Fp8Config(), _quantize_fp8_halves(w1, w2) |
| 749 | |
| 750 | if quantization == "fp8_blocked": |
| 751 | block_shape = [128, 128] |
| 752 | return Fp8Config(weight_block_size=block_shape), _quantize_fp8_halves( |
| 753 | w1, w2, block_shape |
| 754 | ) |
| 755 | |
| 756 | if quantization == "modelopt_fp8": |
| 757 | qw = _quantize_fp8_halves(w1, w2) |
| 758 | # why? |
| 759 | qw.w13_input_scale = torch.ones( |
| 760 | num_experts, dtype=torch.float32, device=w1.device |
| 761 | ) |
| 762 | # why? |
| 763 | qw.w2_input_scale = torch.ones( |
| 764 | num_experts, dtype=torch.float32, device=w2.device |
| 765 | ) |
| 766 | quant_config = ModelOptFp8Config( |
| 767 | quant_method="FP8", |
| 768 | is_checkpoint_fp8_serialized=True, |
| 769 | kv_cache_quant_method=None, |
| 770 | exclude_modules=[], |
| 771 | ) |
| 772 | return quant_config, qw |
| 773 | |
| 774 | if quantization == "modelopt_fp4": |
| 775 | # Quantize full w13 at once so both gate/up halves share the same |
| 776 | # global scale per expert. process_weights_after_loading uses |
| 777 | # w13_weight_scale_2[:, 0] for the entire tensor, so the two shard |
| 778 | # scales must match. |
| 779 | w1q, w1s, w1gs = moe_quantize_weights(w1, None, "nvfp4", False, None) |
| 780 | assert w1s is not None and w1gs is not None |
| 781 | |
| 782 | w2q, w2s, w2gs = moe_quantize_weights(w2, None, "nvfp4", False, None) |
| 783 | assert w2s is not None and w2gs is not None |
| 784 | |
| 785 | qw = QuantizedWeights( |
| 786 | w13_weight=w1q, |
| 787 | w2_weight=w2q, |
| 788 | w13_weight_scale=w1s, |
| 789 | w2_weight_scale=w2s, |
| 790 | # weight_scale_2 = 1/w_gs: the kernel computes |
| 791 | # g_alphas = a_scale * w_scale_2, and correct dequant needs 1/w_gs. |
| 792 | # Expand per-expert scalar to (E, 2) for the two shards. |
| 793 | w13_weight_scale_2=(1.0 / w1gs).unsqueeze(1).expand(-1, 2).contiguous(), |
| 794 | w2_weight_scale_2=1.0 / w2gs, |
| 795 | w13_input_scale=torch.ones( |
| 796 | (num_experts, 2), dtype=torch.float32, device=w1.device |
| 797 | ), |
| 798 | w2_input_scale=torch.ones( |
| 799 | num_experts, dtype=torch.float32, device=w2.device |
| 800 | ), |
| 801 | ) |
| 802 | quant_config = ModelOptNvFp4Config( |
| 803 | is_checkpoint_nvfp4_serialized=True, |
| 804 | kv_cache_quant_algo=None, |
| 805 | exclude_modules=[], |
| 806 | ) |
| 807 | return quant_config, qw |
| 808 | |
| 809 | raise NotImplementedError(f"Unsupported quantization: {quantization}") |
| 810 | |
| 811 | |
| 812 | @dataclass |
| 813 | class SharedExpertsConfig: |
| 814 | w1: torch.Tensor |
| 815 | w2: torch.Tensor |
| 816 | w1_s: torch.Tensor | None = None |
| 817 | w2_s: torch.Tensor | None = None |
| 818 | quant_dtype: torch.dtype | str | None = None |
| 819 | |
| 820 | |
| 821 | @dataclass |
| 822 | class MoETestData: |
| 823 | """Container for MOE test data and transforms.""" |
| 824 | |
| 825 | w1: torch.Tensor |
| 826 | w2: torch.Tensor |
| 827 | hidden_states: torch.Tensor |
| 828 | router_logits: torch.Tensor |
| 829 | shared_experts_config: SharedExpertsConfig | None |
| 830 | gate: torch.nn.Module | None |
| 831 | routed_input_transform: torch.nn.Module | None |
| 832 | routed_output_transform: torch.nn.Module | None |
| 833 | routed_expert_hidden_size: int |
| 834 | |
| 835 | |
| 836 | class SimpleGate(torch.nn.Module): |
| 837 | """Simple gate module for testing: computes router logits from hidden states.""" |
| 838 | |
| 839 | def __init__( |
| 840 | self, |
| 841 | hidden_size: int, |
| 842 | num_experts: int, |
| 843 | dtype: torch.dtype, |
| 844 | device: str = "cuda", |
| 845 | ): |
| 846 | super().__init__() |
| 847 | self.weight = torch.nn.Parameter( |
| 848 | torch.randn(num_experts, hidden_size, device=device, dtype=dtype) / 10 |
| 849 | ) |
| 850 | |
| 851 | def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: |
| 852 | """Returns (router_logits, None) to match expected signature.""" |
| 853 | router_logits = torch.nn.functional.linear(hidden_states, self.weight) |
| 854 | return router_logits, None |
| 855 | |
| 856 | |
| 857 | class SimpleRoutedInputTransform(torch.nn.Module): |
| 858 | """Simple linear transform for testing routed input transform |
| 859 | (e.g., latent projection). |
| 860 | """ |
| 861 | |
| 862 | def __init__( |
| 863 | self, |
| 864 | in_features: int, |
| 865 | out_features: int, |
| 866 | dtype: torch.dtype, |
| 867 | device: str = "cuda", |
| 868 | ): |
| 869 | super().__init__() |
| 870 | self.weight = torch.nn.Parameter( |
| 871 | torch.randn(out_features, in_features, device=device, dtype=dtype) / 10 |
| 872 | ) |
| 873 | |
| 874 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 875 | return torch.nn.functional.linear(x, self.weight) |
| 876 | |
| 877 | |
| 878 | def create_shared_experts_from_config( |
| 879 | shared_experts_config: SharedExpertsConfig | None, |
| 880 | in_dtype: torch.dtype, |
| 881 | tp_size: int = 1, |
| 882 | tp_rank: int = 0, |
| 883 | is_sequence_parallel: bool = False, |
| 884 | device: torch.device | str | None = "cuda", |
| 885 | ) -> TestMLP | None: |
| 886 | """Create TestMLP for shared experts from config. |
| 887 | |
| 888 | Args: |
| 889 | shared_experts_config: Configuration for shared experts |
| 890 | in_dtype: Output data type |
| 891 | tp_size: Tensor parallel size (for weight chunking) |
| 892 | tp_rank: Tensor parallel rank (for weight chunking) |
| 893 | device: Device to move weights to (optional) |
| 894 | |
| 895 | Returns: |
| 896 | TestMLP instance or None if config is None |
| 897 | |
| 898 | """ |
| 899 | if shared_experts_config is None: |
| 900 | return None |
| 901 | |
| 902 | s_w1 = shared_experts_config.w1 |
| 903 | s_w2 = shared_experts_config.w2 |
| 904 | |
| 905 | # Apply TP chunking if needed |
| 906 | if tp_size > 1 and not is_sequence_parallel: |
| 907 | s_w1 = tp_chunk_gate_up(s_w1, tp_rank, tp_size, dim=1, device=device) |
| 908 | s_w2 = chunk_by_rank(s_w2, tp_rank, tp_size, dim=0, device=device) |
| 909 | else: |
| 910 | s_w1 = s_w1.to(device) |
| 911 | s_w2 = s_w2.to(device) |
| 912 | |
| 913 | return TestMLP(w1=s_w1, w2=s_w2, out_dtype=in_dtype) |
| 914 | |
| 915 | |
| 916 | # Make version that takes a MoETestConfig? |
| 917 | def setup_moe_test_data( |
| 918 | m: int, |
| 919 | k: int, |
| 920 | n: int, |
| 921 | num_experts: int, |
| 922 | in_dtype: torch.dtype, |
| 923 | use_shared_experts: bool, |
| 924 | use_gate: bool, |
| 925 | use_routed_input_transform: bool, |
| 926 | backend: str | None, |
| 927 | device: str = "cuda", |
| 928 | ) -> MoETestData: |
| 929 | """Setup test data and transforms for MOE tests. |
| 930 | |
| 931 | Args: |
| 932 | m: Number of tokens |
| 933 | k: Hidden size |
| 934 | n: Intermediate size |
| 935 | num_experts: Number of experts |
| 936 | in_dtype: Data type for tensors |
| 937 | use_shared_experts: Whether to create shared experts config |
| 938 | use_gate: Whether to create gate module |
| 939 | use_routed_input_transform: Whether to create routed input/output transforms |
| 940 | device: Device to create tensors on ("cuda" or "cpu") |
| 941 | |
| 942 | Returns: |
| 943 | MoETestData containing all test data and transforms |
| 944 | |
| 945 | """ |
| 946 | # For latent MoE: latent_size = k // 2 |
| 947 | latent_size = k // 2 |
| 948 | |
| 949 | # k = maybe_roundup_layer_hidden_size(k, in_dtype, backend) |
| 950 | # latent_size = maybe_roundup_layer_hidden_size(latent_size, in_dtype, backend) |
| 951 | |
| 952 | # Determine dimensions for routed experts (may be transformed) |
| 953 | # For latent MoE, routed experts operate entirely in latent space |
| 954 | # (k//2). The routed_output_transform then projects back to k before |
| 955 | # adding with shared experts. |
| 956 | # w1: (E, 2*N, latent_size) - input latent_size |
| 957 | # w2: (E, latent_size, N) - output latent_size (fused_experts returns |
| 958 | # same shape as input) |
| 959 | routed_expert_hidden_size = latent_size if use_routed_input_transform else k |
| 960 | |
| 961 | # Create expert weights |
| 962 | (w1, _, _, _), (w2, _, _, _) = make_test_weights( |
| 963 | num_experts, |
| 964 | n, |
| 965 | routed_expert_hidden_size, # Both w1 input and w2 output use latent_size |
| 966 | in_dtype=in_dtype, |
| 967 | ) |
| 968 | |
| 969 | # Create shared experts config if needed |
| 970 | if use_shared_experts: |
| 971 | shared_experts_config = SharedExpertsConfig( |
| 972 | w1=torch.randn((k, n * 2), device=device, dtype=in_dtype) / 15, |
| 973 | w2=torch.randn((n, k), device=device, dtype=in_dtype) / 15, |
| 974 | ) |
| 975 | else: |
| 976 | shared_experts_config = None |
| 977 | |
| 978 | # Create routed input transform if needed |
| 979 | routed_input_transform = ( |
| 980 | SimpleRoutedInputTransform(k, latent_size, in_dtype, device=device) |
| 981 | if use_routed_input_transform |
| 982 | else None |
| 983 | ) |
| 984 | |
| 985 | # Create gate if needed |
| 986 | # Note: gate is called AFTER routed_input_transform, so it should expect |
| 987 | # the transformed dimension (latent_size) when routed_input_transform is used |
| 988 | gate_input_dim = latent_size if use_routed_input_transform else k |
| 989 | gate = ( |
| 990 | SimpleGate(gate_input_dim, num_experts, in_dtype, device=device) |
| 991 | if use_gate |
| 992 | else None |
| 993 | ) |
| 994 | |
| 995 | # Create routed output transform if needed (projects latent space back to original) |
| 996 | routed_output_transform = ( |
| 997 | SimpleRoutedInputTransform(latent_size, k, in_dtype, device=device) |
| 998 | if use_routed_input_transform |
| 999 | else None |
| 1000 | ) |
| 1001 | |
| 1002 | # Create test inputs |
| 1003 | hidden_states = torch.randn((m, k), device=device, dtype=in_dtype) / 10 |
| 1004 | router_logits = torch.randn((m, num_experts), device=device, dtype=in_dtype) |
| 1005 | |
| 1006 | return MoETestData( |
| 1007 | w1=w1, |
| 1008 | w2=w2, |
| 1009 | hidden_states=hidden_states, |
| 1010 | router_logits=router_logits, |
| 1011 | shared_experts_config=shared_experts_config, |
| 1012 | gate=gate, |
| 1013 | routed_input_transform=routed_input_transform, |
| 1014 | routed_output_transform=routed_output_transform, |
| 1015 | routed_expert_hidden_size=routed_expert_hidden_size, |
| 1016 | ) |
| 1017 | |
| 1018 | |
| 1019 | def make_fused_moe_layer( |
| 1020 | quantization: str | None, |
| 1021 | use_ep: bool, |
| 1022 | hidden_size: int, |
| 1023 | intermediate_size: int, |
| 1024 | in_dtype: torch.dtype, |
| 1025 | tp_size: int, |
| 1026 | ep_size: int, |
| 1027 | dp_size: int, |
| 1028 | w1: torch.Tensor, |
| 1029 | w2: torch.Tensor, |
| 1030 | top_k: int, |
| 1031 | global_num_experts: int, |
| 1032 | renormalize: bool = False, |
| 1033 | shared_experts: torch.nn.Module | None = None, |
| 1034 | use_grouped_topk: bool = False, |
| 1035 | topk_group: int | None = None, |
| 1036 | num_expert_group: int | None = None, |
| 1037 | custom_routing_function: Callable | None = None, |
| 1038 | scoring_func: str = "softmax", |
| 1039 | routed_scaling_factor: float = 1.0, |
| 1040 | e_score_correction_bias: torch.Tensor | None = None, |
| 1041 | apply_router_weight_on_input: bool = False, |
| 1042 | activation: str = "silu", |
| 1043 | indices_type: torch.dtype | None = None, |
| 1044 | expert_map: torch.Tensor | None = None, |
| 1045 | enable_eplb: bool = False, |
| 1046 | expert_load_view: torch.Tensor | None = None, |
| 1047 | logical_to_physical_map: torch.Tensor | None = None, |
| 1048 | logical_replica_count: torch.Tensor | None = None, |
| 1049 | num_redundant_experts: int = 0, |
| 1050 | has_bias: bool = False, |
| 1051 | gate: torch.nn.Module | None = None, |
| 1052 | routed_input_transform: torch.nn.Module | None = None, |
| 1053 | routed_output_transform: torch.nn.Module | None = None, |
| 1054 | pcp_size: int | None = 1, |
| 1055 | is_sequence_parallel: bool = False, |
| 1056 | ) -> MoERunner: |
| 1057 | quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts) |
| 1058 | |
| 1059 | kwargs = dict() |
| 1060 | kwargs["shared_experts"] = shared_experts |
| 1061 | |
| 1062 | # Add gate and routed_input_transform if provided |
| 1063 | if gate is not None: |
| 1064 | kwargs["gate"] = gate |
| 1065 | |
| 1066 | if routed_input_transform is not None: |
| 1067 | kwargs["routed_input_transform"] = routed_input_transform |
| 1068 | kwargs["routed_output_transform"] = routed_output_transform |
| 1069 | |
| 1070 | layer = FusedMoEFactory( |
| 1071 | num_experts=global_num_experts, |
| 1072 | top_k=top_k, |
| 1073 | hidden_size=hidden_size, |
| 1074 | intermediate_size=intermediate_size, |
| 1075 | params_dtype=in_dtype, |
| 1076 | renormalize=renormalize, |
| 1077 | use_grouped_topk=use_grouped_topk, |
| 1078 | num_expert_group=num_expert_group, |
| 1079 | topk_group=topk_group, |
| 1080 | quant_config=quant_config, |
| 1081 | tp_size=tp_size, |
| 1082 | dp_size=dp_size, |
| 1083 | pcp_size=pcp_size, |
| 1084 | prefix="from_forward_context", |
| 1085 | custom_routing_function=custom_routing_function, |
| 1086 | scoring_func=scoring_func, |
| 1087 | routed_scaling_factor=routed_scaling_factor, |
| 1088 | e_score_correction_bias=e_score_correction_bias, |
| 1089 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 1090 | activation=activation, |
| 1091 | enable_eplb=enable_eplb, |
| 1092 | num_redundant_experts=num_redundant_experts, |
| 1093 | has_bias=has_bias, |
| 1094 | is_sequence_parallel=is_sequence_parallel, |
| 1095 | **kwargs, |
| 1096 | ) |
| 1097 | |
| 1098 | weight_scale_name = getattr( |
| 1099 | layer._quant_method, "weight_scale_name", "weight_scale" |
| 1100 | ) |
| 1101 | |
| 1102 | for name, value in [ |
| 1103 | ("w13_weight", qw.w13_weight), |
| 1104 | ("w2_weight", qw.w2_weight), |
| 1105 | (f"w13_{weight_scale_name}", qw.w13_weight_scale), |
| 1106 | (f"w2_{weight_scale_name}", qw.w2_weight_scale), |
| 1107 | ("w13_weight_scale_2", qw.w13_weight_scale_2), |
| 1108 | ("w2_weight_scale_2", qw.w2_weight_scale_2), |
| 1109 | ("w13_input_scale", qw.w13_input_scale), |
| 1110 | ("w2_input_scale", qw.w2_input_scale), |
| 1111 | ]: |
| 1112 | if value is not None: |
| 1113 | layer.routed_experts.register_parameter( |
| 1114 | name, torch.nn.Parameter(value, requires_grad=False) |
| 1115 | ) |
| 1116 | |
| 1117 | layer._quant_method.process_weights_after_loading(layer.routed_experts) |
| 1118 | |
| 1119 | return layer |
| 1120 | |
| 1121 | |
| 1122 | def make_fake_moe_layer( |
| 1123 | w1: torch.Tensor, |
| 1124 | w2: torch.Tensor, |
| 1125 | top_k: int, |
| 1126 | global_num_experts: int, |
| 1127 | in_dtype: torch.dtype, |
| 1128 | quantization: str | None, |
| 1129 | renormalize: bool = False, |
| 1130 | shared_experts_config: SharedExpertsConfig | None = None, |
| 1131 | use_grouped_topk: bool = False, |
| 1132 | topk_group: int | None = None, |
| 1133 | num_expert_group: int | None = None, |
| 1134 | custom_routing_function: Callable | None = None, |
| 1135 | scoring_func: str = "softmax", |
| 1136 | routed_scaling_factor: float = 1.0, |
| 1137 | e_score_correction_bias: torch.Tensor | None = None, |
| 1138 | apply_router_weight_on_input: bool = False, |
| 1139 | activation: str = "silu", |
| 1140 | indices_type: torch.dtype | None = None, |
| 1141 | expert_map: torch.Tensor | None = None, |
| 1142 | expert_load_view: torch.Tensor | None = None, |
| 1143 | logical_to_physical_map: torch.Tensor | None = None, |
| 1144 | logical_replica_count: torch.Tensor | None = None, |
| 1145 | num_redundant_experts: int = 0, |
| 1146 | gate: torch.nn.Module | None = None, |
| 1147 | routed_input_transform: torch.nn.Module | None = None, |
| 1148 | routed_output_transform: torch.nn.Module | None = None, |
| 1149 | use_ep: bool = False, |
| 1150 | tp_size: int = 1, |
| 1151 | dp_size: int = 1, |
| 1152 | ep_size: int = 1, |
| 1153 | is_sequence_parallel: bool = False, |
| 1154 | ) -> Callable: |
| 1155 | quant_dtype = None |
| 1156 | activation = MoEActivation.from_str(activation) |
| 1157 | |
| 1158 | router = create_fused_moe_router( |
| 1159 | top_k=top_k, |
| 1160 | global_num_experts=global_num_experts, |
| 1161 | renormalize=renormalize, |
| 1162 | use_grouped_topk=use_grouped_topk, |
| 1163 | num_expert_group=num_expert_group, |
| 1164 | topk_group=topk_group, |
| 1165 | custom_routing_function=custom_routing_function, |
| 1166 | scoring_func=scoring_func, |
| 1167 | routed_scaling_factor=routed_scaling_factor, |
| 1168 | e_score_correction_bias=e_score_correction_bias, |
| 1169 | num_fused_shared_experts=0, # TODO |
| 1170 | ) |
| 1171 | |
| 1172 | if quant_dtype is not None: |
| 1173 | w1, w1_s, _ = moe_quantize_weights(w1, None, quant_dtype, False, None) |
| 1174 | w2, w2_s, _ = moe_quantize_weights(w2, None, quant_dtype, False, None) |
| 1175 | else: |
| 1176 | w1_s = None |
| 1177 | w2_s = None |
| 1178 | |
| 1179 | shared_experts = create_shared_experts_from_config( |
| 1180 | shared_experts_config, |
| 1181 | in_dtype, |
| 1182 | ) |
| 1183 | |
| 1184 | quant_config = FusedMoEQuantConfig.make( |
| 1185 | quant_dtype, |
| 1186 | w1_scale=w1_s, |
| 1187 | w2_scale=w2_s, |
| 1188 | ) |
| 1189 | |
| 1190 | def _moe( |
| 1191 | hidden_states: torch.Tensor, |
| 1192 | router_logits: torch.Tensor, |
| 1193 | ) -> torch.Tensor: |
| 1194 | # Save original hidden_states for shared experts (before transform) |
| 1195 | original_hidden_states = hidden_states |
| 1196 | |
| 1197 | # Apply routed input transform if provided |
| 1198 | if routed_input_transform is not None: |
| 1199 | hidden_states = routed_input_transform(hidden_states) |
| 1200 | |
| 1201 | # If gate provided, compute router_logits from hidden_states |
| 1202 | # Note: gate operates on transformed hidden_states (after |
| 1203 | # routed_input_transform) |
| 1204 | if gate is not None: |
| 1205 | router_logits, _ = gate(hidden_states) |
| 1206 | |
| 1207 | topk_weights, topk_ids = router.select_experts( |
| 1208 | hidden_states=hidden_states, |
| 1209 | router_logits=router_logits, |
| 1210 | topk_indices_dtype=indices_type, |
| 1211 | ) |
| 1212 | |
| 1213 | # Shared experts use original (untransformed) hidden_states |
| 1214 | if shared_experts is not None: |
| 1215 | shared_output = shared_experts(original_hidden_states) |
| 1216 | else: |
| 1217 | shared_output = None |
| 1218 | |
| 1219 | # Routed experts use transformed hidden_states |
| 1220 | output = fused_experts( |
| 1221 | hidden_states=hidden_states, |
| 1222 | w1=w1, |
| 1223 | w2=w2, |
| 1224 | quant_config=quant_config, |
| 1225 | topk_weights=topk_weights, |
| 1226 | topk_ids=topk_ids, |
| 1227 | activation=activation, |
| 1228 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 1229 | global_num_experts=global_num_experts, |
| 1230 | expert_map=expert_map, |
| 1231 | ) |
| 1232 | |
| 1233 | # Apply routed output transform if provided |
| 1234 | # (e.g., latent space -> original space) |
| 1235 | if routed_output_transform is not None: |
| 1236 | output = routed_output_transform(output) |
| 1237 | |
| 1238 | if shared_experts is not None: |
| 1239 | assert shared_output is not None |
| 1240 | output += shared_output |
| 1241 | |
| 1242 | # Apply TP/DP reduction if not already reduced |
| 1243 | # if (tp_size > 1 or dp_size > 1): |
| 1244 | # output = tensor_model_parallel_all_reduce(output) |
| 1245 | |
| 1246 | return output |
| 1247 | |
| 1248 | return _moe |
| 1249 | |
| 1250 | |
| 1251 | def _test_body_regular( |
| 1252 | moe_layer: MoERunner, |
| 1253 | hidden_states: torch.Tensor, |
| 1254 | router_logits: torch.Tensor, |
| 1255 | vllm_config: VllmConfig, |
| 1256 | num_tokens: int, |
| 1257 | num_tokens_across_dp: torch.Tensor, |
| 1258 | **kwargs, |
| 1259 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1260 | """Regular MoE test body: compare layer output to baseline.""" |
| 1261 | baseline_output = kwargs["baseline_output"] |
| 1262 | |
| 1263 | with set_forward_context( |
| 1264 | None, |
| 1265 | vllm_config, |
| 1266 | num_tokens=num_tokens, |
| 1267 | num_tokens_across_dp=num_tokens_across_dp, |
| 1268 | ): |
| 1269 | output = sp_wrapper(moe_layer)(hidden_states, router_logits) |
| 1270 | |
| 1271 | return baseline_output, output |
| 1272 | |
| 1273 | |
| 1274 | def _test_body_eplb( |
| 1275 | moe_layer: MoERunner, |
| 1276 | hidden_states: torch.Tensor, |
| 1277 | router_logits: torch.Tensor, |
| 1278 | vllm_config: VllmConfig, |
| 1279 | num_tokens: int, |
| 1280 | num_tokens_across_dp: torch.Tensor, |
| 1281 | cpu_group, |
| 1282 | in_dtype: torch.dtype, |
| 1283 | quantization: str | None, |
| 1284 | use_ep: bool, |
| 1285 | tp_size: int, |
| 1286 | ep_size: int, |
| 1287 | dp_size: int, |
| 1288 | w1: torch.Tensor, |
| 1289 | w2: torch.Tensor, |
| 1290 | num_experts: int, |
| 1291 | k: int, |
| 1292 | n: int, |
| 1293 | top_k: int, |
| 1294 | shared_experts, |
| 1295 | gate: torch.nn.Module | None, |
| 1296 | routed_input_transform: torch.nn.Module | None, |
| 1297 | routed_output_transform: torch.nn.Module | None, |
| 1298 | **kwargs, |
| 1299 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1300 | device = torch.accelerator.current_accelerator() |
| 1301 | |
| 1302 | is_sequence_parallel = moe_layer.moe_config.moe_parallel_config.is_sequence_parallel |
| 1303 | |
| 1304 | """EPLB test body: compare output before and after expert weight rearrangement.""" |
| 1305 | # Get "before" output with original weight arrangement |
| 1306 | with set_forward_context( |
| 1307 | None, |
| 1308 | vllm_config, |
| 1309 | num_tokens=num_tokens, |
| 1310 | num_tokens_across_dp=num_tokens_across_dp, |
| 1311 | ): |
| 1312 | output_before = sp_wrapper(moe_layer)(hidden_states, router_logits) |
| 1313 | |
| 1314 | # Create a fresh MoERunner layer with enable_eplb=True |
| 1315 | # Delete the original layer's registration so the constructor can |
| 1316 | # re-use the same "from_forward_context" prefix |
| 1317 | cc = vllm_config.compilation_config |
| 1318 | del cc.static_forward_context["from_forward_context"] |
| 1319 | cc.static_all_moe_layers.remove("from_forward_context") |
| 1320 | |
| 1321 | # Determine hidden size for MoE layer |
| 1322 | # When using routed_input_transform, experts operate in latent space |
| 1323 | hidden_size_for_layer = k // 2 if routed_input_transform is not None else k |
| 1324 | |
| 1325 | eplb_moe_layer = make_fused_moe_layer( |
| 1326 | quantization=quantization, |
| 1327 | use_ep=use_ep, |
| 1328 | hidden_size=hidden_size_for_layer, |
| 1329 | intermediate_size=n, |
| 1330 | in_dtype=in_dtype, |
| 1331 | tp_size=tp_size, |
| 1332 | ep_size=ep_size, |
| 1333 | dp_size=dp_size, |
| 1334 | w1=w1, |
| 1335 | w2=w2, |
| 1336 | top_k=top_k, |
| 1337 | global_num_experts=num_experts, |
| 1338 | shared_experts=shared_experts, |
| 1339 | enable_eplb=True, |
| 1340 | gate=gate, |
| 1341 | routed_input_transform=routed_input_transform, |
| 1342 | routed_output_transform=routed_output_transform, |
| 1343 | is_sequence_parallel=is_sequence_parallel, |
| 1344 | ) |
| 1345 | |
| 1346 | # All ranks must generate the same permutation |
| 1347 | initial_indices = torch.arange(num_experts, dtype=torch.long) |
| 1348 | shuffled_indices = initial_indices[torch.randperm(num_experts)] |
| 1349 | |
| 1350 | expert_weights = [list(eplb_moe_layer.get_expert_weights())] |
| 1351 | |
| 1352 | expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] |
| 1353 | assert vllm_config.parallel_config.eplb_config.communicator is not None, ( |
| 1354 | "EPLB communicator backend must be set by ParallelConfig" |
| 1355 | ) |
| 1356 | communicator = create_eplb_communicator( |
| 1357 | group_coordinator=get_eplb_group(), |
| 1358 | backend=vllm_config.parallel_config.eplb_config.communicator, |
| 1359 | expert_weights=expert_weights, |
| 1360 | expert_buffer=expert_buffer, |
| 1361 | ) |
| 1362 | |
| 1363 | # Rearrange expert weights across EP ranks |
| 1364 | rearrange_expert_weights_inplace( |
| 1365 | old_global_expert_indices=initial_indices.unsqueeze(0), |
| 1366 | new_global_expert_indices=shuffled_indices.unsqueeze(0), |
| 1367 | expert_weights=expert_weights, |
| 1368 | expert_buffer=expert_buffer, |
| 1369 | ep_group=cpu_group, |
| 1370 | communicator=communicator, |
| 1371 | ) |
| 1372 | |
| 1373 | # Build logical_to_physical_map from shuffled_indices |
| 1374 | # shuffled_indices[physical] = logical, we need the inverse |
| 1375 | logical_to_physical = torch.empty(num_experts, dtype=torch.int32, device=device) |
| 1376 | logical_to_physical[shuffled_indices.to(device)] = torch.arange( |
| 1377 | num_experts, dtype=torch.int32, device=device |
| 1378 | ) |
| 1379 | |
| 1380 | eplb_moe_layer.set_eplb_state( |
| 1381 | moe_layer_idx=0, |
| 1382 | expert_load_view=torch.zeros( |
| 1383 | (1, num_experts), |
| 1384 | dtype=torch.int32, |
| 1385 | device=device, |
| 1386 | ), |
| 1387 | logical_to_physical_map=logical_to_physical.reshape(num_experts, 1).unsqueeze( |
| 1388 | 0 |
| 1389 | ), |
| 1390 | logical_replica_count=torch.ones( |
| 1391 | (1, num_experts), |
| 1392 | dtype=torch.int32, |
| 1393 | device=device, |
| 1394 | ), |
| 1395 | ) |
| 1396 | |
| 1397 | eplb_moe_layer.router.eplb_state.should_record_tensor = torch.ones( |
| 1398 | (), dtype=torch.bool, device=device |
| 1399 | ) |
| 1400 | eplb_moe_layer.router.eplb_state.num_unpadded_tokens_tensors = [ |
| 1401 | torch.tensor(0, dtype=torch.int32, device=device) |
| 1402 | ] |
| 1403 | |
| 1404 | # Get "after" output with rearranged weights and EPLB routing |
| 1405 | with set_forward_context( |
| 1406 | None, |
| 1407 | vllm_config, |
| 1408 | num_tokens=num_tokens, |
| 1409 | num_tokens_across_dp=num_tokens_across_dp, |
| 1410 | ): |
| 1411 | output_after = sp_wrapper(eplb_moe_layer)(hidden_states, router_logits) |
| 1412 | |
| 1413 | return output_before, output_after |
| 1414 | |
| 1415 | |
| 1416 | # TODO: make this take a MoETestConfig |
| 1417 | def _run_one_config( |
| 1418 | vllm_config: VllmConfig, |
| 1419 | ep_size: int, # Expert parallel size (total across all ranks) |
| 1420 | dp_size: int, # Data parallel size (number of DP groups) |
| 1421 | tp_size: int, # Tensor parallel size OR sequence parallel size (when use_ep=True) |
| 1422 | dp_rank: int, # Current rank in data parallel dimension |
| 1423 | tp_rank: int, # Current rank in tensor/sequence parallel dimension |
| 1424 | is_sequence_parallel: bool, # Whether to use sequence parallelism |
| 1425 | m: int, |
| 1426 | n: int, |
| 1427 | k: int, |
| 1428 | num_experts: int, |
| 1429 | top_k: int, |
| 1430 | quantization: str | None, |
| 1431 | backend: str | None, |
| 1432 | test_body_fn: Callable, |
| 1433 | use_shared_experts: bool, |
| 1434 | use_gate: bool, |
| 1435 | use_routed_input_transform: bool, |
| 1436 | input_scale: float = 1.0, |
| 1437 | **kwargs, |
| 1438 | ) -> None: |
| 1439 | """Generic test loop that sets up environment and delegates to test_body_fn. |
| 1440 | |
| 1441 | Parameter Interpretation: |
| 1442 | - When is_sequence_parallel=False (standard TP or EP): |
| 1443 | * ep_size: Number of expert parallel ranks (or 1 if no EP) |
| 1444 | * tp_size: Number of tensor parallel ranks (or 1 if no TP) |
| 1445 | * Weights are chunked by ep_size (experts) and tp_size (tensors) |
| 1446 | |
| 1447 | - When is_sequence_parallel=True (EP + sequence splitting): |
| 1448 | * ep_size: Number of expert parallel ranks (equals dp_size * tp_size) |
| 1449 | * tp_size: Number of ranks to split sequence across (becomes sp_size in |
| 1450 | MoERunner) |
| 1451 | * Weights are chunked by ep_size (experts) but NOT by tp_size |
| 1452 | * Input sequences are chunked by tp_size (via sp_wrapper) |
| 1453 | """ |
| 1454 | try: |
| 1455 | set_random_seed(7) |
| 1456 | |
| 1457 | use_ep = ep_size > 1 |
| 1458 | |
| 1459 | assert vllm_config.parallel_config.enable_expert_parallel == use_ep |
| 1460 | |
| 1461 | in_dtype = torch.bfloat16 |
| 1462 | device = torch.accelerator.current_accelerator() |
| 1463 | |
| 1464 | if not is_workspace_manager_initialized(): |
| 1465 | init_workspace_manager(device) |
| 1466 | |
| 1467 | # Create test data and transforms |
| 1468 | test_data = setup_moe_test_data( |
| 1469 | m=m, |
| 1470 | k=k, |
| 1471 | n=n, |
| 1472 | num_experts=num_experts, |
| 1473 | in_dtype=in_dtype, |
| 1474 | use_shared_experts=use_shared_experts, |
| 1475 | use_gate=use_gate, |
| 1476 | use_routed_input_transform=use_routed_input_transform, |
| 1477 | backend=backend, |
| 1478 | device=device, |
| 1479 | ) |
| 1480 | |
| 1481 | # Extract data from test_data |
| 1482 | hidden_states = test_data.hidden_states |
| 1483 | if input_scale != 1.0: |
| 1484 | hidden_states = hidden_states * input_scale |
| 1485 | router_logits = test_data.router_logits |
| 1486 | w1 = test_data.w1 |
| 1487 | w2 = test_data.w2 |
| 1488 | shared_experts_config = test_data.shared_experts_config |
| 1489 | gate = test_data.gate |
| 1490 | routed_input_transform = test_data.routed_input_transform |
| 1491 | routed_output_transform = test_data.routed_output_transform |
| 1492 | activation = "silu" |
| 1493 | |
| 1494 | # Create baseline layer with FULL weights (no EP chunking) |
| 1495 | # Baseline represents the expected output using full model |
| 1496 | baseline_layer = make_fake_moe_layer( |
| 1497 | w1=w1, |
| 1498 | w2=w2, |
| 1499 | top_k=top_k, |
| 1500 | global_num_experts=num_experts, |
| 1501 | in_dtype=in_dtype, |
| 1502 | quantization=quantization, |
| 1503 | renormalize=False, |
| 1504 | shared_experts_config=shared_experts_config, |
| 1505 | gate=gate, |
| 1506 | routed_input_transform=routed_input_transform, |
| 1507 | routed_output_transform=routed_output_transform, |
| 1508 | use_ep=use_ep, |
| 1509 | tp_size=tp_size, |
| 1510 | ep_size=ep_size, |
| 1511 | dp_size=dp_size, |
| 1512 | activation=activation, |
| 1513 | is_sequence_parallel=is_sequence_parallel, |
| 1514 | ) |
| 1515 | |
| 1516 | with set_current_vllm_config(vllm_config): |
| 1517 | # Compute baseline output with SP wrapper if needed |
| 1518 | # sp_wrapper handles sequence chunking/gathering for SP |
| 1519 | baseline_output = sp_wrapper(baseline_layer, is_sequence_parallel)( |
| 1520 | hidden_states, router_logits |
| 1521 | ) |
| 1522 | |
| 1523 | del baseline_layer |
| 1524 | torch.accelerator.empty_cache() |
| 1525 | |
| 1526 | with set_current_vllm_config(vllm_config): |
| 1527 | # Chunk weights for EP BEFORE creating MoERunner. |
| 1528 | # MoERunner uses EP-chunked weights and handles reductions internally. |
| 1529 | if ep_size > 1: |
| 1530 | # Split experts across ranks (dimension 0 is the expert dimension) |
| 1531 | # When EP is enabled, use EP group rank and ep_size for chunking |
| 1532 | ep_rank = get_ep_group().rank_in_group |
| 1533 | w1 = chunk_by_rank(w1, ep_rank, ep_size, dim=0, device=device) |
| 1534 | w2 = chunk_by_rank(w2, ep_rank, ep_size, dim=0, device=device) |
| 1535 | |
| 1536 | # Chunk weights for TP (only if NOT doing sequence parallelism) |
| 1537 | # Sequence parallelism splits tokens/sequences, not weight tensors |
| 1538 | if tp_size > 1 and not is_sequence_parallel: |
| 1539 | w1 = tp_chunk_gate_up(w1, tp_rank, tp_size, dim=1, device=device) |
| 1540 | w2 = chunk_by_rank(w2, tp_rank, tp_size, dim=2, device=device) |
| 1541 | |
| 1542 | # Setup shared experts if needed |
| 1543 | # In SP mode, shared experts should NOT be TP-chunked (same as routed |
| 1544 | # experts). |
| 1545 | # tp_size is used for sequence splitting, not weight splitting |
| 1546 | shared_experts = create_shared_experts_from_config( |
| 1547 | shared_experts_config, |
| 1548 | in_dtype, |
| 1549 | tp_size, |
| 1550 | tp_rank, |
| 1551 | is_sequence_parallel, |
| 1552 | device, |
| 1553 | ) |
| 1554 | |
| 1555 | # Determine hidden size for MoE layer |
| 1556 | # When using routed_input_transform, experts operate in latent space |
| 1557 | hidden_size_for_layer = k // 2 if routed_input_transform is not None else k |
| 1558 | |
| 1559 | # Create initial MoE layer |
| 1560 | moe_layer = make_fused_moe_layer( |
| 1561 | quantization=quantization, |
| 1562 | use_ep=use_ep, |
| 1563 | hidden_size=hidden_size_for_layer, |
| 1564 | intermediate_size=n, |
| 1565 | in_dtype=in_dtype, |
| 1566 | tp_size=tp_size, |
| 1567 | ep_size=ep_size, |
| 1568 | dp_size=dp_size, |
| 1569 | w1=w1, |
| 1570 | w2=w2, |
| 1571 | top_k=top_k, |
| 1572 | global_num_experts=num_experts, |
| 1573 | shared_experts=shared_experts, |
| 1574 | gate=gate, |
| 1575 | routed_input_transform=routed_input_transform, |
| 1576 | routed_output_transform=routed_output_transform, |
| 1577 | activation=activation, |
| 1578 | is_sequence_parallel=is_sequence_parallel, |
| 1579 | ) |
| 1580 | |
| 1581 | num_tokens = m |
| 1582 | # num_tokens_across_dp should have one entry per DP group, not per |
| 1583 | # total rank. |
| 1584 | # When EP is enabled, dp_size represents the number of DP groups |
| 1585 | num_tokens_across_dp = torch.tensor( |
| 1586 | [num_tokens] * dp_size, |
| 1587 | device=device, |
| 1588 | dtype=torch.int, |
| 1589 | ) |
| 1590 | |
| 1591 | # Call the test body function with all necessary context |
| 1592 | expected, actual = test_body_fn( |
| 1593 | moe_layer=moe_layer, |
| 1594 | hidden_states=hidden_states, |
| 1595 | router_logits=router_logits, |
| 1596 | vllm_config=vllm_config, |
| 1597 | num_tokens=num_tokens, |
| 1598 | num_tokens_across_dp=num_tokens_across_dp, |
| 1599 | in_dtype=in_dtype, |
| 1600 | quantization=quantization, |
| 1601 | use_ep=use_ep, |
| 1602 | tp_size=tp_size, |
| 1603 | ep_size=ep_size, |
| 1604 | dp_size=dp_size, |
| 1605 | w1=w1, |
| 1606 | w2=w2, |
| 1607 | num_experts=num_experts, |
| 1608 | k=k, |
| 1609 | n=n, |
| 1610 | m=m, |
| 1611 | top_k=top_k, |
| 1612 | shared_experts=shared_experts, |
| 1613 | gate=gate, |
| 1614 | routed_input_transform=routed_input_transform, |
| 1615 | routed_output_transform=routed_output_transform, |
| 1616 | baseline_output=baseline_output, |
| 1617 | **kwargs, |
| 1618 | ) |
| 1619 | |
| 1620 | # Common tolerance logic |
| 1621 | # TODO: consider associating tolerances with quant methods. |
| 1622 | if quantization is None: |
| 1623 | if k >= 2048: |
| 1624 | atol, rtol = 7.6e-2, 7.6e-2 |
| 1625 | else: |
| 1626 | atol, rtol = 3.5e-2, 3.5e-2 |
| 1627 | elif quantization in ("fp8", "fp8_blocked", "modelopt_fp8"): |
| 1628 | atol, rtol = 6.5e-2, 6.5e-2 |
| 1629 | elif quantization == "modelopt_fp4": |
| 1630 | if k >= 2048: |
| 1631 | atol = rtol = 1e-1 + (k * 1e-4) |
| 1632 | else: |
| 1633 | atol = rtol = 1e-1 |
| 1634 | |
| 1635 | if backend == "allgather_reducescatter" and tp_size > 1: |
| 1636 | atol += 2e-1 |
| 1637 | rtol += 2e-1 |
| 1638 | else: |
| 1639 | atol, rtol = 6e-2, 6e-2 |
| 1640 | |
| 1641 | torch.testing.assert_close(expected, actual, atol=atol, rtol=rtol) |
| 1642 | finally: |
| 1643 | torch.accelerator.synchronize() |
| 1644 | |
| 1645 | |
| 1646 | # Test for non-parallel cases (world_size == 1) - backend doesn't matter |
| 1647 | @pytest.mark.parametrize("m, n, k", SHAPE_COMBOS) |
| 1648 | @pytest.mark.parametrize("num_experts", NUM_EXPERTS) |
| 1649 | @pytest.mark.parametrize("top_k", TOP_KS) |
| 1650 | @pytest.mark.parametrize("quantization", QUANT_METHODS) |
| 1651 | @pytest.mark.parametrize("use_shared_experts", [False, True]) |
| 1652 | @pytest.mark.parametrize("use_gate", [False, True]) |
| 1653 | @pytest.mark.parametrize("use_routed_input_transform", [False, True]) |
| 1654 | def test_moe_layer_no_parallel( |
| 1655 | m: int, |
| 1656 | n: int, |
| 1657 | k: int, |
| 1658 | num_experts: int, |
| 1659 | top_k: int, |
| 1660 | quantization: str | None, |
| 1661 | use_shared_experts: bool, |
| 1662 | use_gate: bool, |
| 1663 | use_routed_input_transform: bool, |
| 1664 | monkeypatch, |
| 1665 | ): |
| 1666 | """Test MoE layer without parallelism (dp_size=1, tp_size=1, use_ep=False).""" |
| 1667 | if os.environ.get("VLLM_LOGGING_LEVEL") is None: |
| 1668 | monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR") |
| 1669 | |
| 1670 | # Needed since weights will already be in e4m3fnuz format and the |
| 1671 | # normalize_e4m3fn_to_e4m3fnuz() function is not being tested here. |
| 1672 | if current_platform.is_fp8_fnuz(): |
| 1673 | override_normalize_e4m3fn_to_e4m3fnuz() |
| 1674 | |
| 1675 | test_config = MoETestConfig( |
| 1676 | m, |
| 1677 | n, |
| 1678 | k, |
| 1679 | num_experts, |
| 1680 | top_k, |
| 1681 | torch.bfloat16, |
| 1682 | quantization, |
| 1683 | use_shared_experts, |
| 1684 | use_gate, |
| 1685 | use_routed_input_transform, |
| 1686 | ) |
| 1687 | |
| 1688 | valid, reason = is_valid_config(test_config) |
| 1689 | if not valid: |
| 1690 | pytest.skip(reason) |
| 1691 | |
| 1692 | set_random_seed(7) |
| 1693 | |
| 1694 | parallel_config = ParallelConfig() |
| 1695 | compilation_config = CompilationConfig() |
| 1696 | compilation_config.pass_config.fuse_allreduce_rms = False |
| 1697 | |
| 1698 | vllm_config = VllmConfig( |
| 1699 | parallel_config=parallel_config, compilation_config=compilation_config |
| 1700 | ) |
| 1701 | |
| 1702 | # Initialize distributed environment for single GPU |
| 1703 | _set_vllm_config(vllm_config, 1, rank=0, local_rank=0) |
| 1704 | |
| 1705 | _run_one_config( |
| 1706 | vllm_config, |
| 1707 | test_config.ep_size, |
| 1708 | test_config.dp_size, |
| 1709 | test_config.tp_size, |
| 1710 | 0, |
| 1711 | 0, |
| 1712 | False, |
| 1713 | test_config.m, |
| 1714 | test_config.n, |
| 1715 | test_config.k, |
| 1716 | test_config.num_experts, |
| 1717 | test_config.top_k, |
| 1718 | test_config.quantization, |
| 1719 | test_config.backend, |
| 1720 | _test_body_regular, |
| 1721 | use_shared_experts=test_config.use_shared_experts, |
| 1722 | use_gate=test_config.use_gate, |
| 1723 | use_routed_input_transform=test_config.use_routed_input_transform, |
| 1724 | ) |
| 1725 | |
| 1726 | |
| 1727 | def _test_body_config(test_config: MoETestConfig, cpu_group, **kwargs): |
| 1728 | if not test_config.enable_eplb: |
| 1729 | return _test_body_regular(**kwargs) |
| 1730 | else: |
| 1731 | return _test_body_eplb(**kwargs, cpu_group=cpu_group) |
| 1732 | |
| 1733 | |
| 1734 | def _parallel_worker( |
| 1735 | pgi: ProcessGroupInfo, |
| 1736 | vllm_config: VllmConfig, |
| 1737 | cpu_group, |
| 1738 | test_configs: list[MoETestConfig], |
| 1739 | verbosity: int, |
| 1740 | failure_report_path: str | None = None, |
| 1741 | deep_ep_handle_keepalive: list[object] | None = None, |
| 1742 | **kwargs, |
| 1743 | ) -> None: |
| 1744 | set_random_seed(7) |
| 1745 | is_logging_rank = pgi.rank == 0 |
| 1746 | |
| 1747 | total = 0 |
| 1748 | passed = 0 |
| 1749 | failed = 0 |
| 1750 | failure_details = [] |
| 1751 | |
| 1752 | dp_rank = vllm_config.parallel_config.data_parallel_rank |
| 1753 | |
| 1754 | if current_platform.is_fp8_fnuz(): |
| 1755 | override_normalize_e4m3fn_to_e4m3fnuz() |
| 1756 | |
| 1757 | for test_config in test_configs: |
| 1758 | cc = vllm_config.compilation_config |
| 1759 | if "from_forward_context" in cc.static_forward_context: |
| 1760 | del cc.static_forward_context["from_forward_context"] |
| 1761 | cc.static_all_moe_layers.remove("from_forward_context") |
| 1762 | |
| 1763 | tp_rank = pgi.rank % test_config.tp_size |
| 1764 | |
| 1765 | if verbosity > 0 and is_logging_rank: |
| 1766 | print(f"subtest: {test_config.id()}", end="") |
| 1767 | |
| 1768 | local_failed = False |
| 1769 | local_error: str | None = None |
| 1770 | try: |
| 1771 | _run_one_config( |
| 1772 | vllm_config, |
| 1773 | test_config.ep_size, |
| 1774 | test_config.dp_size, |
| 1775 | test_config.tp_size, |
| 1776 | dp_rank, |
| 1777 | tp_rank, |
| 1778 | test_config.is_sequence_parallel, |
| 1779 | test_config.m, |
| 1780 | test_config.n, |
| 1781 | test_config.k, |
| 1782 | test_config.num_experts, |
| 1783 | test_config.top_k, |
| 1784 | test_config.quantization, |
| 1785 | test_config.backend, |
| 1786 | functools.partial( |
| 1787 | _test_body_config, test_config=test_config, cpu_group=cpu_group |
| 1788 | ), |
| 1789 | use_shared_experts=test_config.use_shared_experts, |
| 1790 | use_gate=test_config.use_gate, |
| 1791 | use_routed_input_transform=test_config.use_routed_input_transform, |
| 1792 | ) |
| 1793 | except Exception: |
| 1794 | local_failed = True |
| 1795 | local_error = traceback.format_exc() |
| 1796 | finally: |
| 1797 | # DeepEP managers are not reliably reusable across many subtests in |
| 1798 | # a single worker process. Tear them down after each DeepEP case so |
| 1799 | # later subtests do not inherit stale communication state. Skip this |
| 1800 | # on ROCm: rocSHMEM cannot reinitialize the allocator after a DeepEP |
| 1801 | # buffer is destroyed in the same process. |
| 1802 | if current_platform.is_cuda() and test_config.backend in DEEPEP_BACKENDS: |
| 1803 | torch.accelerator.synchronize() |
| 1804 | all2all_manager = get_ep_group().device_communicator.all2all_manager |
| 1805 | if all2all_manager is not None: |
| 1806 | all2all_manager.destroy() |
| 1807 | elif ( |
| 1808 | deep_ep_handle_keepalive is not None |
| 1809 | and current_platform.is_rocm() |
| 1810 | and test_config.backend in DEEPEP_BACKENDS |
| 1811 | ): |
| 1812 | # The manager cache is weak. Keep its handle alive until the |
| 1813 | # launcher hard-exits this ROCm worker, otherwise each subtest |
| 1814 | # implicitly destroys and recreates the DeepEP buffer. |
| 1815 | all2all_manager = get_ep_group().device_communicator.all2all_manager |
| 1816 | if all2all_manager is not None: |
| 1817 | handle_cache = getattr(all2all_manager, "handle_cache", None) |
| 1818 | if handle_cache is not None and not deep_ep_handle_keepalive: |
| 1819 | with handle_cache._lock: |
| 1820 | cached_handles = list(handle_cache._cache.values()) |
| 1821 | if cached_handles: |
| 1822 | deep_ep_handle_keepalive.extend(cached_handles) |
| 1823 | total = total + 1 |
| 1824 | torch.distributed.barrier() |
| 1825 | |
| 1826 | any_failed_tensor = torch.tensor( |
| 1827 | [int(local_failed)], device=pgi.device, dtype=torch.int32 |
| 1828 | ) |
| 1829 | torch.distributed.all_reduce( |
| 1830 | any_failed_tensor, op=torch.distributed.ReduceOp.MAX |
| 1831 | ) |
| 1832 | any_failed = bool(any_failed_tensor.item()) |
| 1833 | |
| 1834 | if any_failed: |
| 1835 | failed = failed + 1 |
| 1836 | |
| 1837 | gathered_errors = [None] * pgi.world_size |
| 1838 | torch.distributed.all_gather_object( |
| 1839 | gathered_errors, local_error, group=cpu_group |
| 1840 | ) |
| 1841 | first_error = next( |
| 1842 | (error for error in gathered_errors if error is not None), |
| 1843 | "unknown distributed failure", |
| 1844 | ) |
| 1845 | assert first_error is not None |
| 1846 | failure_details.append(f"{test_config.id()}\n{first_error.rstrip()}") |
| 1847 | |
| 1848 | if verbosity > 0 and is_logging_rank: |
| 1849 | print(" FAILED") |
| 1850 | print(first_error.rstrip()) |
| 1851 | elif is_logging_rank: |
| 1852 | print("F", end="") |
| 1853 | else: |
| 1854 | passed = passed + 1 |
| 1855 | if verbosity > 0 and is_logging_rank: |
| 1856 | print(" PASSED") |
| 1857 | elif is_logging_rank: |
| 1858 | print(".", end="") |
| 1859 | |
| 1860 | skipped = total - (passed + failed) |
| 1861 | |
| 1862 | fails = f"{failed} failed" if failed > 0 else "" |
| 1863 | sep = ", " if fails != "" else "" |
| 1864 | skips = f"{sep}{skipped} skipped" if skipped > 0 else "" |
| 1865 | sep = ", " if skips != "" or fails != "" else "" |
| 1866 | passes = f"{sep}{passed} passed" if passed > 0 else "" |
| 1867 | |
| 1868 | report = ( |
| 1869 | f"============= {fails}{skips}{passes} of {total} total subcases =============" |
| 1870 | ) |
| 1871 | |
| 1872 | if is_logging_rank: |
| 1873 | sep = "\n" if verbosity == 0 else "" |
| 1874 | print(f"{sep}{report}") |
| 1875 | |
| 1876 | if failed > 0: |
| 1877 | failure_details_str = "\n\n".join(failure_details) |
| 1878 | failure_report = ( |
| 1879 | f"\n============= Failed subcases =============\n" |
| 1880 | f"{failure_details_str}\n{report}" |
| 1881 | ) |
| 1882 | if is_logging_rank and failure_report_path is not None: |
| 1883 | with open(failure_report_path, "a", encoding="utf-8") as report_file: |
| 1884 | report_file.write(failure_report) |
| 1885 | if is_logging_rank: |
| 1886 | raise RuntimeError(failure_report) |
| 1887 | |
| 1888 | |
| 1889 | def _parallel_worker_rocm_deepep( |
| 1890 | pgi: ProcessGroupInfo, |
| 1891 | vllm_config: VllmConfig, |
| 1892 | cpu_group, |
| 1893 | test_configs: list[MoETestConfig], |
| 1894 | verbosity: int, |
| 1895 | failure_report_path: str | None = None, |
| 1896 | deep_ep_handle_keepalive: list[object] | None = None, |
| 1897 | **kwargs, |
| 1898 | ) -> None: |
| 1899 | """Run a ROCm DeepEP batch without unsafe Python/HIP teardown.""" |
| 1900 | assert current_platform.is_rocm() |
| 1901 | assert deep_ep_handle_keepalive is not None |
| 1902 | |
| 1903 | exit_code = 1 |
| 1904 | try: |
| 1905 | _parallel_worker( |
| 1906 | pgi, |
| 1907 | vllm_config, |
| 1908 | cpu_group, |
| 1909 | test_configs, |
| 1910 | verbosity, |
| 1911 | failure_report_path=failure_report_path, |
| 1912 | deep_ep_handle_keepalive=deep_ep_handle_keepalive, |
| 1913 | **kwargs, |
| 1914 | ) |
| 1915 | exit_code = 0 |
| 1916 | except BaseException as ex: |
| 1917 | print(ex) |
| 1918 | traceback.print_exc() |
| 1919 | finally: |
| 1920 | try: |
| 1921 | torch.accelerator.synchronize() |
| 1922 | # Do not run vLLM cleanup: it explicitly destroys the DeepEP |
| 1923 | # buffer. Destroy only the default group, matching the accepted |
| 1924 | # ROCm workaround in tests/kernels/moe/parallel_utils.py. |
| 1925 | torch.distributed.destroy_process_group() |
| 1926 | except BaseException: |
| 1927 | traceback.print_exc() |
| 1928 | exit_code = 1 |
| 1929 | finally: |
| 1930 | # Bypass the HIP atexit use-after-free fixed upstream by |
| 1931 | # https://github.com/ROCm/rocm-systems/pull/6942. |
| 1932 | sys.stdout.flush() |
| 1933 | sys.stderr.flush() |
| 1934 | os._exit(exit_code) |
| 1935 | |
| 1936 | |
| 1937 | def _test_body_mori_graph( |
| 1938 | moe_layer: MoERunner, |
| 1939 | hidden_states: torch.Tensor, |
| 1940 | router_logits: torch.Tensor, |
| 1941 | vllm_config: VllmConfig, |
| 1942 | num_tokens: int, |
| 1943 | num_tokens_across_dp: torch.Tensor, |
| 1944 | baseline_output: torch.Tensor, |
| 1945 | **kwargs, |
| 1946 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1947 | from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( |
| 1948 | AiterExperts, |
| 1949 | ) |
| 1950 | from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import ( |
| 1951 | MoriPrepareAndFinalize, |
| 1952 | ) |
| 1953 | |
| 1954 | kernel = moe_layer._quant_method.moe_kernel |
| 1955 | assert kernel is not None |
| 1956 | assert isinstance(kernel.prepare_finalize, MoriPrepareAndFinalize) |
| 1957 | assert isinstance(kernel.fused_experts, AiterExperts) |
| 1958 | assert not kernel.prepare_finalize.use_fp8_dispatch |
| 1959 | |
| 1960 | def check(actual: torch.Tensor, expected: torch.Tensor) -> None: |
| 1961 | torch.testing.assert_close(actual, expected, atol=3.5e-2, rtol=3.5e-2) |
| 1962 | # These inputs produce small outputs; an absolute-only budget can |
| 1963 | # accept all-zero output. Also bound error relative to the reference. |
| 1964 | ref_norm = expected.float().norm() |
| 1965 | assert ref_norm > 0 |
| 1966 | assert (actual.float() - expected.float()).norm() < 3.5e-2 * ref_norm |
| 1967 | |
| 1968 | # Forward-context metadata belongs on CPU, outside graph capture. |
| 1969 | with set_forward_context( |
| 1970 | None, |
| 1971 | vllm_config, |
| 1972 | num_tokens=num_tokens, |
| 1973 | num_tokens_across_dp=num_tokens_across_dp.cpu(), |
| 1974 | ): |
| 1975 | initial_hidden = hidden_states.clone() |
| 1976 | initial_router = router_logits.clone() |
| 1977 | for _ in range(2): |
| 1978 | eager_output = moe_layer(hidden_states, router_logits) |
| 1979 | saved_eager_output = eager_output.clone() |
| 1980 | torch.accelerator.synchronize() |
| 1981 | |
| 1982 | graph = torch.cuda.CUDAGraph() |
| 1983 | with torch.cuda.graph(graph): |
| 1984 | captured_output = moe_layer(hidden_states, router_logits) |
| 1985 | graph.replay() |
| 1986 | original_output = captured_output.clone() |
| 1987 | check(original_output, baseline_output) |
| 1988 | |
| 1989 | for scale, shift in [(-0.5, 1), (1.75, 3)]: |
| 1990 | hidden_states.copy_(initial_hidden * scale) |
| 1991 | router_logits.copy_(initial_router.roll(shift, dims=-1)) |
| 1992 | expected = moe_layer(hidden_states, router_logits).clone() |
| 1993 | graph.replay() |
| 1994 | check(captured_output, expected) |
| 1995 | assert (expected.float() - original_output.float()).norm() > ( |
| 1996 | 0.1 * original_output.float().norm() |
| 1997 | ) |
| 1998 | |
| 1999 | # Later calls must not overwrite an earlier eager return tensor. |
| 2000 | check(eager_output, saved_eager_output) |
| 2001 | check(original_output, baseline_output) |
| 2002 | return baseline_output, original_output |
| 2003 | |
| 2004 | |
| 2005 | def _mori_graph_worker( |
| 2006 | pgi: ProcessGroupInfo, |
| 2007 | vllm_config: VllmConfig, |
| 2008 | cpu_group, |
| 2009 | m: int, |
| 2010 | n: int, |
| 2011 | ) -> None: |
| 2012 | from vllm._aiter_ops import rocm_aiter_ops |
| 2013 | |
| 2014 | rocm_aiter_ops.refresh_env_variables() |
| 2015 | _run_one_config( |
| 2016 | vllm_config=vllm_config, |
| 2017 | ep_size=2, |
| 2018 | dp_size=2, |
| 2019 | tp_size=1, |
| 2020 | dp_rank=pgi.rank, |
| 2021 | tp_rank=0, |
| 2022 | is_sequence_parallel=False, |
| 2023 | m=m, |
| 2024 | n=n, |
| 2025 | k=512, |
| 2026 | num_experts=8, |
| 2027 | top_k=2, |
| 2028 | quantization=None, |
| 2029 | backend="mori_high_throughput", |
| 2030 | test_body_fn=_test_body_mori_graph, |
| 2031 | use_shared_experts=False, |
| 2032 | use_gate=False, |
| 2033 | use_routed_input_transform=False, |
| 2034 | input_scale=1.0 + pgi.rank, |
| 2035 | ) |
| 2036 | |
| 2037 | |
| 2038 | @pytest.mark.parametrize("m,n", [(1, 256), (32, 512), (45, 256)]) |
| 2039 | def test_moe_layer_mori_graph(m: int, n: int, monkeypatch): |
| 2040 | """BF16 EP dispatch/combine preserves routing and outputs across replays.""" |
| 2041 | from vllm._aiter_ops import is_aiter_found_and_supported |
| 2042 | |
| 2043 | if not current_platform.is_rocm() or not has_mori(): |
| 2044 | pytest.skip("Requires ROCm and MoRI") |
| 2045 | if not is_aiter_found_and_supported(): |
| 2046 | pytest.skip("Requires supported AITER kernels") |
| 2047 | if current_platform.device_count() < 2: |
| 2048 | pytest.skip("Requires two GPUs") |
| 2049 | |
| 2050 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 2051 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") |
| 2052 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0") |
| 2053 | compilation_config = CompilationConfig() |
| 2054 | compilation_config.pass_config.fuse_allreduce_rms = False |
| 2055 | vllm_config = VllmConfig( |
| 2056 | parallel_config=ParallelConfig( |
| 2057 | data_parallel_size=2, |
| 2058 | tensor_parallel_size=1, |
| 2059 | enable_expert_parallel=True, |
| 2060 | all2all_backend="mori_high_throughput", |
| 2061 | ), |
| 2062 | compilation_config=compilation_config, |
| 2063 | scheduler_config=SchedulerConfig.default_factory( |
| 2064 | max_num_batched_tokens=64, max_num_seqs=64 |
| 2065 | ), |
| 2066 | ) |
| 2067 | parallel_launch_with_config(2, _mori_graph_worker, vllm_config, None, m, n) |
| 2068 | |
| 2069 | |
| 2070 | # TODO: add torch.compile tests and graph coverage for other backends. |
| 2071 | @pytest.mark.parametrize("dp_size, tp_size, use_ep", PARALLEL_COMBOS) |
| 2072 | @pytest.mark.parametrize("backend", BACKENDS) |
| 2073 | @pytest.mark.parametrize("enable_eplb", [False, True]) |
| 2074 | def test_moe_layer( |
| 2075 | dp_size: int, |
| 2076 | tp_size: int, |
| 2077 | use_ep: bool, |
| 2078 | backend: str, |
| 2079 | enable_eplb: bool, |
| 2080 | monkeypatch, |
| 2081 | pytestconfig, |
| 2082 | subtests, |
| 2083 | ): |
| 2084 | """Test MoE layer with parallelism (multi-GPU or TP/EP enabled). |
| 2085 | |
| 2086 | For non-parallel cases (world_size == 1), use test_moe_layer_no_parallel instead. |
| 2087 | """ |
| 2088 | num_gpus = current_platform.device_count() |
| 2089 | world_size = tp_size * dp_size |
| 2090 | # When use_ep=True: FusedMoEParallelConfig flattens tp_size across dp ranks |
| 2091 | # Result: ep_size = dp_size * pcp_size * tp_size |
| 2092 | # Since pcp_size=1 in these tests: ep_size = dp_size * tp_size = world_size |
| 2093 | # When use_ep=False: no expert parallelism, ep_size = 1 |
| 2094 | ep_size = 1 if not use_ep else world_size |
| 2095 | |
| 2096 | assert world_size > 1 |
| 2097 | |
| 2098 | # Check if enough GPUs available |
| 2099 | if world_size is not None and num_gpus is not None and world_size > num_gpus: |
| 2100 | pytest.skip(f"Not enough GPUs got {num_gpus}, expected {world_size}.") |
| 2101 | |
| 2102 | if enable_eplb and not use_ep: |
| 2103 | pytest.skip("EPLB requires EP.") |
| 2104 | |
| 2105 | if backend in DEEPEP_BACKENDS and not visible_devices_have_peer_access(world_size): |
| 2106 | pytest.skip("DeepEP backends require peer access between visible GPUs.") |
| 2107 | |
| 2108 | verbosity = pytestconfig.getoption("verbose") |
| 2109 | |
| 2110 | if os.environ.get("VLLM_LOGGING_LEVEL") is None: |
| 2111 | monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR") |
| 2112 | |
| 2113 | if ( |
| 2114 | backend in MORI_BACKENDS |
| 2115 | and os.environ.get("VLLM_TEST_ENABLE_MORI_MOE_LAYER") == "1" |
| 2116 | ): |
| 2117 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 2118 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") |
| 2119 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0") |
| 2120 | from vllm._aiter_ops import rocm_aiter_ops |
| 2121 | |
| 2122 | rocm_aiter_ops.refresh_env_variables() |
| 2123 | |
| 2124 | # TODO: cover FlashInfer MoE backends via moe_backend, e.g. |
| 2125 | # moe_backend=flashinfer_trtllm / flashinfer_cutlass / flashinfer_cutedsl |
| 2126 | # (BF16, FP8 and NVFP4 paths), and VLLM_USE_FLASHINFER_MOE_INT4=1. |
| 2127 | |
| 2128 | # Repeated NIXL memory registration in this broad layer matrix fails on |
| 2129 | # gfx950. NIXL EPLB is covered by dedicated tests, so use Gloo here to |
| 2130 | # preserve the layer/EPLB coverage. |
| 2131 | eplb_config = EPLBConfig( |
| 2132 | communicator="torch_gloo" if enable_eplb and on_gfx950() else None |
| 2133 | ) |
| 2134 | parallel_config = ParallelConfig( |
| 2135 | pipeline_parallel_size=1, |
| 2136 | data_parallel_size=dp_size, |
| 2137 | tensor_parallel_size=tp_size, |
| 2138 | enable_expert_parallel=use_ep, |
| 2139 | all2all_backend=backend, |
| 2140 | enable_eplb=enable_eplb, |
| 2141 | eplb_config=eplb_config, |
| 2142 | ) |
| 2143 | |
| 2144 | compilation_config = CompilationConfig() |
| 2145 | # compilation_config.mode = CompilationMode.NONE # for now |
| 2146 | compilation_config.pass_config.fuse_allreduce_rms = False # for now |
| 2147 | |
| 2148 | vllm_config = VllmConfig( |
| 2149 | parallel_config=parallel_config, |
| 2150 | compilation_config=compilation_config, |
| 2151 | scheduler_config=SchedulerConfig.default_factory( |
| 2152 | max_num_batched_tokens=next_power_of_2(MAX_M) |
| 2153 | ), |
| 2154 | ) |
| 2155 | |
| 2156 | test_configs = generate_valid_test_configs( |
| 2157 | backend, ep_size, dp_size, tp_size, enable_eplb, verbosity |
| 2158 | ) |
| 2159 | |
| 2160 | if subtests is not None: |
| 2161 | new_test_configs = [] |
| 2162 | for subtest in subtests.split(","): |
| 2163 | sub_test_config = MoETestConfig.from_id(subtest) |
| 2164 | if sub_test_config in test_configs: |
| 2165 | new_test_configs.append(sub_test_config) |
| 2166 | else: |
| 2167 | pytest.skip( |
| 2168 | f"subtest config {subtest} does not match any valid test " |
| 2169 | "configuration" |
| 2170 | ) |
| 2171 | test_configs = new_test_configs |
| 2172 | |
| 2173 | if len(test_configs) == 0: |
| 2174 | pytest.skip("No supported configs found for this testpoint.") |
| 2175 | |
| 2176 | with tempfile.NamedTemporaryFile( |
| 2177 | prefix="moe-layer-failures-", delete=False |
| 2178 | ) as failure_report_file: |
| 2179 | failure_report_path = failure_report_file.name |
| 2180 | |
| 2181 | test_config_batches = [test_configs] |
| 2182 | if current_platform.is_rocm() and backend == "deepep_low_latency": |
| 2183 | # rocSHMEM cannot destroy one low-latency buffer and initialize a |
| 2184 | # differently sized one in the same process. Use one worker lifetime |
| 2185 | # per compatible buffer shape while preserving every matrix case. |
| 2186 | test_config_batches = _group_deepep_ll_configs(test_configs) |
| 2187 | |
| 2188 | rocm_deepep = current_platform.is_rocm() and backend in DEEPEP_BACKENDS |
| 2189 | parallel_worker = _parallel_worker_rocm_deepep if rocm_deepep else _parallel_worker |
| 2190 | launch_failures: list[str] = [] |
| 2191 | |
| 2192 | try: |
| 2193 | for test_config_batch in test_config_batches: |
| 2194 | report_size_before = os.path.getsize(failure_report_path) |
| 2195 | try: |
| 2196 | parallel_launch_with_config( |
| 2197 | world_size, |
| 2198 | parallel_worker, |
| 2199 | vllm_config, |
| 2200 | None, |
| 2201 | test_config_batch, |
| 2202 | verbosity, |
| 2203 | failure_report_path=failure_report_path, |
| 2204 | deep_ep_handle_keepalive=[] if rocm_deepep else None, |
| 2205 | ) |
| 2206 | except Exception as ex: |
| 2207 | # Normal subtest failures are already in the shared report. |
| 2208 | # Preserve launcher/setup errors that occur before that write. |
| 2209 | if os.path.getsize(failure_report_path) == report_size_before: |
| 2210 | launch_failures.append(str(ex)) |
| 2211 | |
| 2212 | if os.path.getsize(failure_report_path) > 0: |
| 2213 | with open(failure_report_path, encoding="utf-8") as report_file: |
| 2214 | launch_failures.insert(0, report_file.read()) |
| 2215 | if launch_failures: |
| 2216 | pytest.fail("\n\n".join(launch_failures)) |
| 2217 | finally: |
| 2218 | with suppress(FileNotFoundError): |
| 2219 | os.remove(failure_report_path) |
| 2220 | torch.accelerator.synchronize() # TODO: Is this needed? |
| 2221 | torch.accelerator.empty_cache() |