| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for FusedMoEFactory with routed_input_transform. |
| 4 | |
| 5 | Verifies that applying routed_input_transform inside FusedMoEFactory |
| 6 | produces the same results as applying the transform manually outside. |
| 7 | """ |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | |
| 13 | from vllm.config import VllmConfig, set_current_vllm_config |
| 14 | from vllm.forward_context import set_forward_context |
| 15 | from vllm.model_executor.layers.fused_moe import FusedMoEFactory |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.utils.torch_utils import is_torch_equal_or_newer, set_random_seed |
| 18 | |
| 19 | |
| 20 | class SimpleLinear(nn.Module): |
| 21 | """A simple linear transform mimicking latent projection in latent MoE.""" |
| 22 | |
| 23 | def __init__(self, in_features: int, out_features: int, dtype: torch.dtype): |
| 24 | super().__init__() |
| 25 | self.weight = nn.Parameter( |
| 26 | torch.randn(out_features, in_features, device="cuda", dtype=dtype) / 10 |
| 27 | ) |
| 28 | |
| 29 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 30 | return nn.functional.linear(x, self.weight) |
| 31 | |
| 32 | |
| 33 | class SimpleSharedExperts(nn.Module): |
| 34 | """A simple 2-layer MLP mimicking shared experts.""" |
| 35 | |
| 36 | def __init__(self, hidden_size: int, intermediate_size: int, dtype: torch.dtype): |
| 37 | super().__init__() |
| 38 | self.up = nn.Linear( |
| 39 | hidden_size, intermediate_size * 2, bias=False, device="cuda", dtype=dtype |
| 40 | ) |
| 41 | self.down = nn.Linear( |
| 42 | intermediate_size, hidden_size, bias=False, device="cuda", dtype=dtype |
| 43 | ) |
| 44 | with torch.no_grad(): |
| 45 | self.up.weight.div_(10) |
| 46 | self.down.weight.div_(10) |
| 47 | |
| 48 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 49 | gate_up = self.up(x) |
| 50 | gate, up = gate_up.chunk(2, dim=-1) |
| 51 | return self.down(nn.functional.silu(gate) * up) |
| 52 | |
| 53 | |
| 54 | def _assert_close( |
| 55 | actual: torch.Tensor, |
| 56 | expected: torch.Tensor, |
| 57 | atol: float, |
| 58 | rtol: float, |
| 59 | label: str, |
| 60 | ) -> None: |
| 61 | """assert_close that prints diff diagnostics on both success and failure.""" |
| 62 | actual_nans = int(actual.isnan().sum().item()) |
| 63 | expected_nans = int(expected.isnan().sum().item()) |
| 64 | actual_zeros = int((actual == 0).sum().item()) |
| 65 | expected_zeros = int((expected == 0).sum().item()) |
| 66 | n_total = actual.numel() |
| 67 | |
| 68 | diff = (actual - expected).abs() |
| 69 | max_diff = diff.max().item() |
| 70 | mean_diff = diff.mean().item() |
| 71 | n_exceed = int((diff > atol).sum().item()) |
| 72 | pct_exceed = n_exceed / n_total * 100 |
| 73 | |
| 74 | print( |
| 75 | f"[{label}] " |
| 76 | f"shape={list(actual.shape)}, " |
| 77 | f"max_diff={max_diff:.6e}, " |
| 78 | f"mean_diff={mean_diff:.6e}, " |
| 79 | f"exceed_atol({atol})={n_exceed}/{n_total} ({pct_exceed:.2f}%), " |
| 80 | f"actual=[{actual.min().item():.4f}, {actual.max().item():.4f}], " |
| 81 | f"expected=[{expected.min().item():.4f}, {expected.max().item():.4f}], " |
| 82 | f"nan(actual/expected)={actual_nans}/{expected_nans}, " |
| 83 | f"zeros(actual/expected)={actual_zeros}/{expected_zeros}" |
| 84 | ) |
| 85 | |
| 86 | assert actual_nans == 0, ( |
| 87 | f"{label}: actual has {actual_nans}/{n_total} NaN values " |
| 88 | f"(expected has {expected_nans}). " |
| 89 | f"This indicates a kernel bug, not a precision issue." |
| 90 | ) |
| 91 | assert expected_nans == 0, ( |
| 92 | f"{label}: expected has {expected_nans}/{n_total} NaN values. " |
| 93 | f"This indicates a kernel bug, not a precision issue." |
| 94 | ) |
| 95 | |
| 96 | torch.testing.assert_close( |
| 97 | actual, |
| 98 | expected, |
| 99 | atol=atol, |
| 100 | rtol=rtol, |
| 101 | msg=( |
| 102 | f"{label}: max_diff={max_diff:.6e}, mean_diff={mean_diff:.6e}, " |
| 103 | f"exceed_atol({atol})={n_exceed}/{n_total} ({pct_exceed:.2f}%)" |
| 104 | ), |
| 105 | ) |
| 106 | |
| 107 | |
| 108 | @pytest.fixture(autouse=True) |
| 109 | def setup_cuda(): |
| 110 | if not torch.cuda.is_available(): |
| 111 | pytest.skip("CUDA not available") |
| 112 | torch.set_default_device("cuda") |
| 113 | |
| 114 | |
| 115 | @pytest.mark.parametrize("num_tokens", [1, 32]) |
| 116 | @pytest.mark.parametrize("hidden_size,latent_size", [(256, 128), (128, 64)]) |
| 117 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 118 | @pytest.mark.parametrize( |
| 119 | "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False] |
| 120 | ) |
| 121 | @pytest.mark.skipif( |
| 122 | is_torch_equal_or_newer("2.10.0"), |
| 123 | reason="Test fails with PyTorch 2.10.0 see: https://github.com/vllm-project/vllm/issues/33995", |
| 124 | ) |
| 125 | def test_routed_input_transform_inside_vs_outside( |
| 126 | num_tokens: int, |
| 127 | hidden_size: int, |
| 128 | latent_size: int, |
| 129 | dtype: torch.dtype, |
| 130 | use_rocm_aiter: bool, |
| 131 | dist_init, |
| 132 | workspace_init, |
| 133 | monkeypatch, |
| 134 | ): |
| 135 | """Compare FusedMoEFactory with transform inside vs manually applying outside. |
| 136 | Method A (inside): FusedMoEFactory with routed_input_transform |
| 137 | Method B (outside): Manually transform, then FusedMoEFactory without transform |
| 138 | """ |
| 139 | if current_platform.is_rocm(): |
| 140 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0") |
| 141 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1" if use_rocm_aiter else "0") |
| 142 | from vllm._aiter_ops import rocm_aiter_ops |
| 143 | |
| 144 | rocm_aiter_ops.refresh_env_variables() |
| 145 | |
| 146 | set_random_seed(42) |
| 147 | |
| 148 | num_experts = 8 |
| 149 | top_k = 2 |
| 150 | intermediate_size = hidden_size * 2 |
| 151 | |
| 152 | vllm_config = VllmConfig() |
| 153 | vllm_config.compilation_config.static_forward_context = dict() |
| 154 | |
| 155 | shared_experts = SimpleSharedExperts(hidden_size, intermediate_size, dtype) |
| 156 | routed_transform = SimpleLinear(hidden_size, latent_size, dtype) |
| 157 | |
| 158 | with set_current_vllm_config(vllm_config): |
| 159 | # Method A: FusedMoEFactory WITH routed_input_transform |
| 160 | moe_with_transform = FusedMoEFactory( |
| 161 | shared_experts=shared_experts, |
| 162 | routed_input_transform=routed_transform, |
| 163 | num_experts=num_experts, |
| 164 | top_k=top_k, |
| 165 | hidden_size=latent_size, |
| 166 | intermediate_size=intermediate_size, |
| 167 | renormalize=True, |
| 168 | params_dtype=dtype, |
| 169 | tp_size=1, |
| 170 | dp_size=1, |
| 171 | pcp_size=1, |
| 172 | prefix="moe_with_transform", |
| 173 | ) |
| 174 | |
| 175 | # Method B: FusedMoEFactory WITHOUT routed_input_transform |
| 176 | # Note: shared_experts=None because when transform is done outside, |
| 177 | moe_without_transform = FusedMoEFactory( |
| 178 | shared_experts=None, |
| 179 | routed_input_transform=None, |
| 180 | num_experts=num_experts, |
| 181 | top_k=top_k, |
| 182 | hidden_size=latent_size, |
| 183 | intermediate_size=intermediate_size, |
| 184 | renormalize=True, |
| 185 | params_dtype=dtype, |
| 186 | tp_size=1, |
| 187 | dp_size=1, |
| 188 | pcp_size=1, |
| 189 | prefix="moe_without_transform", |
| 190 | ) |
| 191 | |
| 192 | # Weights are created via torch.empty (uninitialized). |
| 193 | # Initialize with seeded random values for reproducibility. |
| 194 | with torch.no_grad(): |
| 195 | moe_with_transform.w13_weight.normal_() |
| 196 | moe_with_transform.w13_weight.div_(10) |
| 197 | moe_with_transform.w2_weight.normal_() |
| 198 | moe_with_transform.w2_weight.div_(10) |
| 199 | moe_without_transform.w13_weight.copy_(moe_with_transform.w13_weight) |
| 200 | moe_without_transform.w2_weight.copy_(moe_with_transform.w2_weight) |
| 201 | |
| 202 | moe_with_transform.quant_method.process_weights_after_loading( |
| 203 | moe_with_transform |
| 204 | ) |
| 205 | moe_without_transform.quant_method.process_weights_after_loading( |
| 206 | moe_without_transform |
| 207 | ) |
| 208 | |
| 209 | hidden_states = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype) |
| 210 | router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype) |
| 211 | |
| 212 | with set_forward_context(None, vllm_config, num_tokens=num_tokens): |
| 213 | # Method A: combined output (shared + routed) |
| 214 | combined_A = moe_with_transform(hidden_states, router_logits) |
| 215 | |
| 216 | # Method B: manually transform, get routed output, add shared |
| 217 | transformed_hidden = routed_transform(hidden_states) |
| 218 | routed_out_B = moe_without_transform(transformed_hidden, router_logits) |
| 219 | shared_out_B = shared_experts(hidden_states) |
| 220 | combined_B = shared_out_B + routed_out_B |
| 221 | |
| 222 | torch.testing.assert_close( |
| 223 | combined_A, |
| 224 | combined_B, |
| 225 | atol=1e-3, |
| 226 | rtol=1e-3, |
| 227 | msg="Combined output should match: transform inside vs outside", |
| 228 | ) |