| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Two-rank MoRI/AITER numerical counterpart to the DeepEP v2 MoE matrix.""" |
| 4 | |
| 5 | from dataclasses import replace |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from tests.kernels.moe.modular_kernel_tools.parallel_utils import ( |
| 11 | parallel_launch_with_config, |
| 12 | ) |
| 13 | from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights |
| 14 | from tests.kernels.utils import torch_experts |
| 15 | from tests.utils import multi_gpu_test |
| 16 | from vllm.config import ParallelConfig, SchedulerConfig, VllmConfig |
| 17 | from vllm.forward_context import set_forward_context |
| 18 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 19 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 20 | maybe_make_prepare_finalize, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.config import ( |
| 23 | FusedMoEParallelConfig, |
| 24 | FusedMoEQuantConfig, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel |
| 27 | from vllm.platforms import current_platform |
| 28 | from vllm.utils.import_utils import has_mori |
| 29 | from vllm.utils.torch_utils import set_random_seed |
| 30 | from vllm.v1.worker.workspace import init_workspace_manager |
| 31 | |
| 32 | MNKS = [ |
| 33 | (1, 256, 256), |
| 34 | (2, 256, 512), |
| 35 | (3, 1024, 2048), |
| 36 | (32, 256, 1024), |
| 37 | (45, 512, 2048), |
| 38 | (64, 1024, 1024), |
| 39 | (222, 1024, 2048), |
| 40 | ] |
| 41 | |
| 42 | |
| 43 | def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph): |
| 44 | from vllm._aiter_ops import rocm_aiter_ops |
| 45 | from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts |
| 46 | from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import ( |
| 47 | MoriPrepareAndFinalize, |
| 48 | ) |
| 49 | |
| 50 | rocm_aiter_ops.refresh_env_variables() |
| 51 | init_workspace_manager(pgi.device) |
| 52 | set_random_seed(47) |
| 53 | dtype = current_platform.fp8_dtype() if fp8 else None |
| 54 | (_, w1, s1, _), (_, w2, s2, _) = make_test_weights( |
| 55 | 32, n, k, quant_dtype=dtype, per_out_ch_quant=fp8 |
| 56 | ) |
| 57 | # Share global expert weights, but give each rank different token inputs. |
| 58 | set_random_seed(101 + pgi.rank) |
| 59 | x = torch.randn(m, k, device=pgi.device, dtype=torch.bfloat16) / 10 |
| 60 | routing = torch.randn(m, 32, device=pgi.device) |
| 61 | # Guarantee a route to both ranks, including the one-token case. |
| 62 | routing[:, 0] = routing[:, 16] = float("inf") |
| 63 | ids = routing.topk(6, dim=-1).indices.to(torch.int32) |
| 64 | weights = torch.randn(m, 6, device=pgi.device, dtype=torch.float32) |
| 65 | parallel = replace( |
| 66 | FusedMoEParallelConfig.make_no_parallel(), |
| 67 | dp_size=2, |
| 68 | dp_rank=pgi.rank, |
| 69 | ep_size=2, |
| 70 | ep_rank=pgi.rank, |
| 71 | use_ep=True, |
| 72 | all2all_backend="mori_high_throughput", |
| 73 | ) |
| 74 | config = replace( |
| 75 | make_dummy_moe_config( |
| 76 | num_experts=32, |
| 77 | num_local_experts=16, |
| 78 | experts_per_token=6, |
| 79 | hidden_dim=k, |
| 80 | intermediate_size=n, |
| 81 | max_num_tokens=256, |
| 82 | ), |
| 83 | moe_parallel_config=parallel, |
| 84 | ) |
| 85 | start, end = pgi.rank * 16, (pgi.rank + 1) * 16 |
| 86 | quant = FusedMoEQuantConfig.make( |
| 87 | dtype, |
| 88 | per_act_token_quant=fp8, |
| 89 | per_out_ch_quant=fp8, |
| 90 | w1_scale=None if s1 is None else s1[start:end].contiguous(), |
| 91 | w2_scale=None if s2 is None else s2[start:end].contiguous(), |
| 92 | ) |
| 93 | prepare = maybe_make_prepare_finalize(config, quant) |
| 94 | assert isinstance(prepare, MoriPrepareAndFinalize) |
| 95 | assert prepare.use_fp8_dispatch == fp8 |
| 96 | assert prepare.num_dispatchers() == 2 |
| 97 | experts = AiterExperts(config, quant) |
| 98 | assert not experts.expects_unquantized_inputs |
| 99 | kernel = FusedMoEKernel(prepare, experts) |
| 100 | local_w1, local_w2 = rocm_aiter_ops.shuffle_weights( |
| 101 | w1[start:end].contiguous(), w2[start:end].contiguous() |
| 102 | ) |
| 103 | local_w1.is_shuffled = local_w2.is_shuffled = True |
| 104 | expert_mask = torch.zeros(32, dtype=torch.int32, device=pgi.device) |
| 105 | expert_mask[start:end] = 1 |
| 106 | |
| 107 | def reference(route_weights): |
| 108 | return torch_experts( |
| 109 | x, |
| 110 | w1, |
| 111 | w2, |
| 112 | route_weights, |
| 113 | ids, |
| 114 | w1_scale=s1, |
| 115 | w2_scale=s2, |
| 116 | quant_dtype=dtype, |
| 117 | per_act_token_quant=fp8, |
| 118 | ) |
| 119 | |
| 120 | def run(): |
| 121 | with set_forward_context(None, vllm_config): |
| 122 | return kernel.apply( |
| 123 | x, |
| 124 | local_w1, |
| 125 | local_w2, |
| 126 | weights, |
| 127 | ids, |
| 128 | activation=MoEActivation.SILU, |
| 129 | global_num_experts=32, |
| 130 | expert_map=expert_mask, |
| 131 | apply_router_weight_on_input=False, |
| 132 | ) |
| 133 | |
| 134 | def check(actual): |
| 135 | expected = reference(weights) |
| 136 | assert actual.shape == expected.shape |
| 137 | assert actual.dtype == expected.dtype |
| 138 | local_weights = weights * ((ids >= start) & (ids < end)) |
| 139 | missing_remote = reference(local_weights) |
| 140 | norm = expected.float().norm() |
| 141 | assert norm > 0 |
| 142 | norm_budget = 0.2 if fp8 else 0.06 |
| 143 | assert (actual.float() - expected.float()).norm() / norm < norm_budget |
| 144 | # The norm check also rejects low-signal zero or rank-local results. |
| 145 | for incorrect in (torch.zeros_like(expected), missing_remote): |
| 146 | assert (incorrect.float() - expected.float()).norm() / norm > norm_budget |
| 147 | if fp8: |
| 148 | assert ( |
| 149 | torch.isclose(actual, expected, atol=0.2, rtol=0.2) |
| 150 | .float() |
| 151 | .mean() |
| 152 | .item() |
| 153 | > 0.99 |
| 154 | ) |
| 155 | else: |
| 156 | torch.testing.assert_close(actual, expected, atol=0.06, rtol=0.06) |
| 157 | assert torch.isfinite(actual).all() |
| 158 | |
| 159 | actual = run() |
| 160 | check(actual) |
| 161 | if graph: |
| 162 | saved = actual.clone() |
| 163 | run() |
| 164 | capture = torch.cuda.CUDAGraph() |
| 165 | with torch.cuda.graph(capture): |
| 166 | captured = run() |
| 167 | for factor in (-0.5, 1.5): |
| 168 | x.mul_(factor) |
| 169 | weights.mul_(-0.75) |
| 170 | capture.replay() |
| 171 | torch.accelerator.synchronize() |
| 172 | check(captured) |
| 173 | torch.testing.assert_close(actual, saved, atol=0, rtol=0) |
| 174 | |
| 175 | |
| 176 | @pytest.mark.skipif( |
| 177 | not current_platform.is_rocm() or not has_mori(), reason="Requires ROCm MoRI" |
| 178 | ) |
| 179 | @pytest.mark.parametrize("m,n,k", MNKS) |
| 180 | @pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8_token_channel"]) |
| 181 | @multi_gpu_test(num_gpus=2) |
| 182 | def test_mori_moe(m, n, k, fp8): |
| 183 | config = VllmConfig( |
| 184 | parallel_config=ParallelConfig( |
| 185 | data_parallel_size=2, |
| 186 | enable_expert_parallel=True, |
| 187 | all2all_backend="mori_high_throughput", |
| 188 | ), |
| 189 | scheduler_config=SchedulerConfig.default_factory(max_num_batched_tokens=256), |
| 190 | ) |
| 191 | env = { |
| 192 | "VLLM_ROCM_USE_AITER": "1", |
| 193 | "VLLM_ROCM_USE_AITER_MOE": "1", |
| 194 | "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS": "0", |
| 195 | } |
| 196 | parallel_launch_with_config(2, _mori_worker, config, env, m, n, k, fp8, m == 32) |