| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for FusedMoEFactory with zero experts. |
| 4 | |
| 5 | Verifies that: |
| 6 | - The ZeroExpertRouter is properly created and used as the layer router. |
| 7 | - A forward pass through FusedMoEFactory with zero experts produces correct output. |
| 8 | - The output decomposes correctly into real expert + zero expert contributions. |
| 9 | |
| 10 | Note: tests generated with Claude. |
| 11 | """ |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | |
| 16 | from vllm.config import VllmConfig, set_current_vllm_config |
| 17 | from vllm.forward_context import get_forward_context, set_forward_context |
| 18 | from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory |
| 19 | from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( |
| 20 | ZeroExpertRouter, |
| 21 | ) |
| 22 | from vllm.v1.worker.workspace import init_workspace_manager |
| 23 | |
| 24 | |
| 25 | @pytest.fixture |
| 26 | def zero_expert_moe(dist_init, default_vllm_config): |
| 27 | """Create a FusedMoEFactory layer with zero experts.""" |
| 28 | num_experts = 4 |
| 29 | top_k = 2 |
| 30 | # hidden_size must be >= 256 for the zero expert identity kernel to |
| 31 | # produce output (its BLOCK_SIZE=256 causes grid=0 when hidden_dim<256). |
| 32 | hidden_size = 256 |
| 33 | intermediate_size = 512 |
| 34 | zero_expert_num = 1 |
| 35 | |
| 36 | e_score_correction_bias = torch.zeros( |
| 37 | num_experts + zero_expert_num, |
| 38 | dtype=torch.float32, |
| 39 | device="cuda", |
| 40 | ) |
| 41 | |
| 42 | vllm_config = VllmConfig() |
| 43 | vllm_config.compilation_config.static_forward_context = dict() |
| 44 | |
| 45 | with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): |
| 46 | init_workspace_manager(torch.accelerator.current_device_index()) |
| 47 | |
| 48 | layer = FusedMoEFactory( |
| 49 | zero_expert_type="identity", |
| 50 | e_score_correction_bias=e_score_correction_bias, |
| 51 | num_experts=num_experts, |
| 52 | top_k=top_k, |
| 53 | hidden_size=hidden_size, |
| 54 | intermediate_size=intermediate_size, |
| 55 | params_dtype=torch.bfloat16, |
| 56 | prefix="test_zero_expert_moe", |
| 57 | renormalize=False, |
| 58 | routed_scaling_factor=1.0, |
| 59 | scoring_func="softmax", |
| 60 | ).cuda() |
| 61 | |
| 62 | layer._quant_method.process_weights_after_loading(layer.routed_experts) |
| 63 | |
| 64 | yield layer, vllm_config |
| 65 | |
| 66 | |
| 67 | @pytest.mark.parametrize("num_tokens", [1, 32]) |
| 68 | def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens): |
| 69 | """Verify that FusedMoEFactory with zero_expert_type creates a ZeroExpertRouter.""" |
| 70 | layer, _ = zero_expert_moe |
| 71 | assert isinstance(layer.router, ZeroExpertRouter), ( |
| 72 | f"Expected ZeroExpertRouter but got {type(layer.router).__name__}." |
| 73 | ) |
| 74 | |
| 75 | |
| 76 | # @pytest.mark.parametrize("num_tokens", [1, 32]) |
| 77 | # def test_zero_expert_moe_no_custom_routing_fn(zero_expert_moe, num_tokens): |
| 78 | # """Verify that custom_routing_function is not set (routing is handled |
| 79 | # by ZeroExpertRouter, not a memoizing closure).""" |
| 80 | # layer, _ = zero_expert_moe |
| 81 | # #assert layer.custom_routing_function is None |
| 82 | |
| 83 | |
| 84 | @pytest.mark.parametrize("num_tokens", [1, 32]) |
| 85 | def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): |
| 86 | """Run a forward pass through FusedMoEFactory with zero experts |
| 87 | and verify output shape. |
| 88 | """ |
| 89 | layer, vllm_config = zero_expert_moe |
| 90 | |
| 91 | hidden_size = layer.routed_experts.hidden_size |
| 92 | num_experts = 4 |
| 93 | zero_expert_num = 1 |
| 94 | total_experts = num_experts + zero_expert_num |
| 95 | |
| 96 | hidden_states = torch.randn( |
| 97 | num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda" |
| 98 | ) |
| 99 | router_logits = torch.randn( |
| 100 | num_tokens, total_experts, dtype=torch.float32, device="cuda" |
| 101 | ) |
| 102 | |
| 103 | # Initialize weights to small random values to avoid NaN from |
| 104 | # uninitialized memory. |
| 105 | with torch.no_grad(): |
| 106 | for param in layer.parameters(): |
| 107 | if param.dtype.is_floating_point: |
| 108 | param.normal_(0, 0.01) |
| 109 | |
| 110 | with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): |
| 111 | get_forward_context().all_moe_layers = None |
| 112 | output = layer.forward(hidden_states, router_logits) |
| 113 | |
| 114 | assert output.shape == hidden_states.shape, ( |
| 115 | f"Expected output shape {hidden_states.shape}, got {output.shape}" |
| 116 | ) |
| 117 | assert output.dtype == hidden_states.dtype |
| 118 | assert not torch.isnan(output).any(), "Output contains NaN values" |
| 119 | |
| 120 | |
| 121 | @pytest.mark.parametrize("num_tokens", [1, 32]) |
| 122 | def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): |
| 123 | """Validate that the FusedMoEFactory output equals a plain FusedMoEFactory |
| 124 | output (real experts only) plus the zero expert contribution. |
| 125 | |
| 126 | The key invariant is: |
| 127 | zero_layer.forward(h, r_full) == plain_layer.forward(h, r_real) |
| 128 | + zero_expert_output |
| 129 | |
| 130 | We create a plain FusedMoEFactory layer with the same weights and real-expert-only |
| 131 | router logits, compute the zero expert output via the ZeroExpertRouter, and |
| 132 | verify the sum matches the FusedMoEFactory output. |
| 133 | """ |
| 134 | layer, vllm_config = zero_expert_moe |
| 135 | num_experts = 4 |
| 136 | zero_expert_num = 1 |
| 137 | total_experts = num_experts + zero_expert_num |
| 138 | |
| 139 | hidden_states = torch.randn( |
| 140 | num_tokens, |
| 141 | layer.routed_experts.hidden_size, |
| 142 | dtype=torch.bfloat16, |
| 143 | device="cuda", |
| 144 | ) |
| 145 | router_logits = torch.randn( |
| 146 | num_tokens, total_experts, dtype=torch.float32, device="cuda" |
| 147 | ) |
| 148 | |
| 149 | with torch.no_grad(): |
| 150 | for param in layer.parameters(): |
| 151 | if param.dtype.is_floating_point: |
| 152 | param.normal_(0, 0.01) |
| 153 | |
| 154 | with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): |
| 155 | get_forward_context().all_moe_layers = None |
| 156 | |
| 157 | # Create a plain FusedMoEFactory layer with the same config but no zero |
| 158 | # experts. Use a separate prefix to avoid collision. |
| 159 | plain_layer = FusedMoEFactory( |
| 160 | num_experts=num_experts, |
| 161 | top_k=layer.routed_experts.top_k, |
| 162 | hidden_size=layer.routed_experts.hidden_size, |
| 163 | intermediate_size=layer.routed_experts.intermediate_size_per_partition, |
| 164 | params_dtype=torch.bfloat16, |
| 165 | prefix="test_zero_expert_moe_plain", |
| 166 | renormalize=False, |
| 167 | scoring_func="softmax", |
| 168 | e_score_correction_bias=layer.routed_experts.e_score_correction_bias, |
| 169 | ).cuda() |
| 170 | |
| 171 | # Share weights from the zero expert layer. |
| 172 | plain_layer.routed_experts.w13_weight.data.copy_( |
| 173 | layer.routed_experts.w13_weight.data |
| 174 | ) |
| 175 | plain_layer.routed_experts.w2_weight.data.copy_( |
| 176 | layer.routed_experts.w2_weight.data |
| 177 | ) |
| 178 | plain_layer._quant_method.process_weights_after_loading( |
| 179 | plain_layer.routed_experts |
| 180 | ) |
| 181 | |
| 182 | # Compute routing via the ZeroExpertRouter. This produces masked |
| 183 | # topk_weights/topk_ids (zero expert entries have weight=0, id=0) |
| 184 | # and stores zero_expert_output as a side effect. |
| 185 | topk_weights, topk_ids = layer.router.select_experts( |
| 186 | hidden_states, router_logits |
| 187 | ) |
| 188 | zero_output = layer.router.zero_expert_output |
| 189 | |
| 190 | # Compute real expert output using the plain layer with the masked |
| 191 | # routing from the ZeroExpertRouter. |
| 192 | real_output = plain_layer._quant_method.apply( |
| 193 | layer=plain_layer.routed_experts, |
| 194 | x=hidden_states, |
| 195 | topk_weights=topk_weights, |
| 196 | topk_ids=topk_ids, |
| 197 | shared_experts=None, |
| 198 | shared_experts_input=None, |
| 199 | ) |
| 200 | |
| 201 | # Get the combined output from the zero expert layer. |
| 202 | full_output = layer.forward(hidden_states, router_logits) |
| 203 | |
| 204 | assert zero_output is not None, "Zero expert output should not be None" |
| 205 | assert not torch.isnan(real_output).any(), "Real expert output has NaN" |
| 206 | assert not torch.isnan(zero_output).any(), "Zero expert output has NaN" |
| 207 | assert not torch.isnan(full_output).any(), "Full output has NaN" |
| 208 | |
| 209 | expected = real_output + zero_output |
| 210 | torch.testing.assert_close( |
| 211 | full_output, |
| 212 | expected, |
| 213 | atol=4e-3, |
| 214 | rtol=4e-3, |
| 215 | msg="FusedMoEFactory output should equal plain FusedMoEFactory output " |
| 216 | "plus zero expert contribution", |
| 217 | ) |
| 218 | |
| 219 | |
| 220 | @pytest.mark.parametrize("num_tokens", [1, 32]) |
| 221 | def test_zero_expert_moe_zero_expert_is_identity(zero_expert_moe, num_tokens): |
| 222 | """Validate zero expert identity behavior. |
| 223 | |
| 224 | When routing strongly favors the zero expert, its contribution should |
| 225 | be a scaled version of hidden_states (identity operation). We verify |
| 226 | this by manually computing the expected zero expert output from the |
| 227 | routing weights and comparing against what the router produces. |
| 228 | """ |
| 229 | layer, vllm_config = zero_expert_moe |
| 230 | num_experts = 4 |
| 231 | zero_expert_num = 1 |
| 232 | total_experts = num_experts + zero_expert_num |
| 233 | |
| 234 | hidden_states = torch.randn( |
| 235 | num_tokens, |
| 236 | layer.routed_experts.hidden_size, |
| 237 | dtype=torch.bfloat16, |
| 238 | device="cuda", |
| 239 | ) |
| 240 | # Strongly bias toward the zero expert (index 4). |
| 241 | router_logits = torch.full( |
| 242 | (num_tokens, total_experts), -10.0, dtype=torch.float32, device="cuda" |
| 243 | ) |
| 244 | router_logits[:, num_experts] = 10.0 # zero expert gets high logit |
| 245 | |
| 246 | with torch.no_grad(): |
| 247 | for param in layer.parameters(): |
| 248 | if param.dtype.is_floating_point: |
| 249 | param.normal_(0, 0.01) |
| 250 | |
| 251 | with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): |
| 252 | get_forward_context().all_moe_layers = None |
| 253 | |
| 254 | # Run routing to get topk_weights/topk_ids before masking. |
| 255 | from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( |
| 256 | fused_topk_bias, |
| 257 | ) |
| 258 | |
| 259 | topk_weights, topk_ids = fused_topk_bias( |
| 260 | hidden_states=hidden_states, |
| 261 | gating_output=router_logits, |
| 262 | e_score_correction_bias=layer.router.e_score_correction_bias.data, |
| 263 | topk=layer.routed_experts.top_k, |
| 264 | renormalize=layer.router.renormalize, |
| 265 | scoring_func=layer.router.scoring_func, |
| 266 | ) |
| 267 | |
| 268 | # Manually compute expected zero expert identity output: |
| 269 | # For each token, sum routing weights assigned to zero expert slots, |
| 270 | # then multiply by hidden_states. |
| 271 | zero_mask = topk_ids >= num_experts |
| 272 | zero_weight_per_token = (topk_weights * zero_mask.float()).sum( |
| 273 | dim=-1, keepdim=True |
| 274 | ) |
| 275 | expected_zero_output = (hidden_states.float() * zero_weight_per_token).to( |
| 276 | hidden_states.dtype |
| 277 | ) |
| 278 | |
| 279 | # Run routing directly to trigger zero expert computation |
| 280 | # without going through the runner (which consumes the output). |
| 281 | layer.router.select_experts(hidden_states, router_logits) |
| 282 | actual_zero_output = layer.router.zero_expert_output |
| 283 | |
| 284 | assert actual_zero_output is not None |
| 285 | assert zero_mask.any(), ( |
| 286 | "With high zero expert logit, at least some slots should route " |
| 287 | "to the zero expert" |
| 288 | ) |
| 289 | |
| 290 | torch.testing.assert_close( |
| 291 | actual_zero_output, |
| 292 | expected_zero_output, |
| 293 | atol=1e-3, |
| 294 | rtol=1e-3, |
| 295 | msg="Zero expert identity output should equal " |
| 296 | "hidden_states * sum(zero_expert_weights)", |
| 297 | ) |