| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | from unittest.mock import patch |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm._aiter_ops import rocm_aiter_ops |
| 10 | from vllm.distributed.eplb.eplb_state import EplbLayerState |
| 11 | from vllm.model_executor.layers.fused_moe.config import RoutingMethodType |
| 12 | from vllm.model_executor.layers.fused_moe.router.base_router import ( |
| 13 | eplb_map_to_physical_and_record, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( |
| 16 | FusedTopKBiasRouter, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.fused_moe.router.fused_topk_router import ( |
| 19 | FusedTopKRouter, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( |
| 22 | GroupedTopKRouter, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.fused_moe.router.router_factory import ( |
| 25 | create_fused_moe_router, |
| 26 | ) |
| 27 | from vllm.model_executor.models.llama4 import Llama4MoE |
| 28 | from vllm.platforms import current_platform |
| 29 | |
| 30 | |
| 31 | def _is_aiter_capable() -> bool: |
| 32 | """Check if the platform supports AITER (gfx942/gfx950/gfx1250).""" |
| 33 | if not current_platform.is_rocm(): |
| 34 | return False |
| 35 | try: |
| 36 | from vllm.platforms.rocm import get_cdna_version |
| 37 | |
| 38 | return get_cdna_version() > 2 |
| 39 | except ImportError: |
| 40 | return False |
| 41 | |
| 42 | |
| 43 | # Test parameters |
| 44 | MK_S = [(32, 256), (64, 512)] |
| 45 | TOP_KS = [2, 4, 6] |
| 46 | NUM_EXPERTS = [8, 16, 64] |
| 47 | |
| 48 | |
| 49 | def test_degenerate_grouped_config_uses_standard_topk() -> None: |
| 50 | router = create_fused_moe_router( |
| 51 | top_k=4, |
| 52 | global_num_experts=128, |
| 53 | use_grouped_topk=True, |
| 54 | num_expert_group=1, |
| 55 | topk_group=1, |
| 56 | scoring_func="softmax", |
| 57 | renormalize=True, |
| 58 | ) |
| 59 | |
| 60 | assert isinstance(router, FusedTopKRouter) |
| 61 | hidden_states, router_logits = make_test_data(32, 256, 128) |
| 62 | |
| 63 | topk_weights, topk_ids = router.select_experts(hidden_states, router_logits) |
| 64 | baseline_weights, baseline_ids = baseline_fused_topk( |
| 65 | router_logits, |
| 66 | top_k=4, |
| 67 | renormalize=True, |
| 68 | ) |
| 69 | |
| 70 | assert_routing_results_close( |
| 71 | topk_weights, |
| 72 | topk_ids, |
| 73 | baseline_weights, |
| 74 | baseline_ids, |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | def test_multiple_expert_groups_use_grouped_topk() -> None: |
| 79 | router = create_fused_moe_router( |
| 80 | top_k=4, |
| 81 | global_num_experts=128, |
| 82 | use_grouped_topk=True, |
| 83 | num_expert_group=8, |
| 84 | topk_group=4, |
| 85 | scoring_func="softmax", |
| 86 | renormalize=True, |
| 87 | ) |
| 88 | |
| 89 | assert isinstance(router, GroupedTopKRouter) |
| 90 | assert not router.skip_padding |
| 91 | |
| 92 | |
| 93 | def test_grouped_topk_padding_skip_must_be_enabled() -> None: |
| 94 | router = create_fused_moe_router( |
| 95 | top_k=4, |
| 96 | global_num_experts=128, |
| 97 | use_grouped_topk=True, |
| 98 | num_expert_group=8, |
| 99 | topk_group=4, |
| 100 | skip_padding=True, |
| 101 | ) |
| 102 | |
| 103 | assert isinstance(router, GroupedTopKRouter) |
| 104 | assert router.skip_padding |
| 105 | |
| 106 | |
| 107 | def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None: |
| 108 | router = create_fused_moe_router( |
| 109 | top_k=4, |
| 110 | global_num_experts=128, |
| 111 | use_grouped_topk=True, |
| 112 | num_expert_group=1, |
| 113 | topk_group=1, |
| 114 | scoring_func="softmax", |
| 115 | renormalize=True, |
| 116 | e_score_correction_bias=torch.empty(128), |
| 117 | ) |
| 118 | |
| 119 | assert isinstance(router, FusedTopKBiasRouter) |
| 120 | |
| 121 | |
| 122 | def test_degenerate_grouped_config_with_bias_keeps_routed_scale() -> None: |
| 123 | router = create_fused_moe_router( |
| 124 | top_k=4, |
| 125 | global_num_experts=128, |
| 126 | use_grouped_topk=True, |
| 127 | num_expert_group=1, |
| 128 | topk_group=1, |
| 129 | scoring_func="softmax", |
| 130 | renormalize=True, |
| 131 | routed_scaling_factor=1.1, |
| 132 | e_score_correction_bias=torch.empty(128), |
| 133 | ) |
| 134 | |
| 135 | assert isinstance(router, FusedTopKBiasRouter) |
| 136 | assert router.routed_scaling_factor == 1.1 |
| 137 | |
| 138 | |
| 139 | def test_degenerate_deepseek_v3_routing_stays_grouped() -> None: |
| 140 | router = create_fused_moe_router( |
| 141 | top_k=4, |
| 142 | global_num_experts=128, |
| 143 | use_grouped_topk=True, |
| 144 | num_expert_group=1, |
| 145 | topk_group=1, |
| 146 | scoring_func="sigmoid", |
| 147 | renormalize=True, |
| 148 | e_score_correction_bias=torch.empty(128), |
| 149 | ) |
| 150 | |
| 151 | assert isinstance(router, GroupedTopKRouter) |
| 152 | assert router.routing_method_type == RoutingMethodType.DeepSeekV3 |
| 153 | |
| 154 | |
| 155 | def test_single_expert_group_with_non_unit_scale_uses_grouped_topk() -> None: |
| 156 | router = create_fused_moe_router( |
| 157 | top_k=4, |
| 158 | global_num_experts=128, |
| 159 | use_grouped_topk=True, |
| 160 | num_expert_group=1, |
| 161 | topk_group=1, |
| 162 | scoring_func="softmax", |
| 163 | renormalize=True, |
| 164 | routed_scaling_factor=1.1, |
| 165 | ) |
| 166 | |
| 167 | assert isinstance(router, GroupedTopKRouter) |
| 168 | |
| 169 | |
| 170 | def test_sigmoid_bias_routing_with_routed_scale_is_minimax2() -> None: |
| 171 | """FlashInfer's MiniMax2 routing applies routed_scaling_factor, so a |
| 172 | non-unit scale (MiniMax-M3 uses 2.0) must not block fused routing.""" |
| 173 | router = create_fused_moe_router( |
| 174 | top_k=4, |
| 175 | global_num_experts=128, |
| 176 | scoring_func="sigmoid", |
| 177 | renormalize=True, |
| 178 | routed_scaling_factor=2.0, |
| 179 | e_score_correction_bias=torch.empty(128), |
| 180 | ) |
| 181 | |
| 182 | assert isinstance(router, FusedTopKBiasRouter) |
| 183 | assert router.routing_method_type == RoutingMethodType.MiniMax2 |
| 184 | |
| 185 | |
| 186 | def test_zero_expert_routing_is_unspecified() -> None: |
| 187 | """Zero experts are resolved in the router, so kernels with built-in |
| 188 | routing must never be selected for them.""" |
| 189 | router = create_fused_moe_router( |
| 190 | top_k=4, |
| 191 | global_num_experts=128, |
| 192 | scoring_func="sigmoid", |
| 193 | renormalize=True, |
| 194 | e_score_correction_bias=torch.empty(160), |
| 195 | zero_expert_type="identity", |
| 196 | num_logical_experts=160, |
| 197 | ) |
| 198 | |
| 199 | assert router.routing_method_type == RoutingMethodType.Unspecified |
| 200 | |
| 201 | |
| 202 | def setup_eplb_state( |
| 203 | enable_eplb: bool, global_num_experts: int |
| 204 | ) -> EplbLayerState | None: |
| 205 | if not enable_eplb: |
| 206 | return None |
| 207 | |
| 208 | # Initialize EPLB state with proper tensors for testing |
| 209 | # For testing purposes, we use a simple 1:1 mapping (no redundant experts) |
| 210 | # expert_load_view: tracks load on each expert (shape: num_experts) |
| 211 | expert_load_view = torch.zeros(global_num_experts, dtype=torch.int32, device="cuda") |
| 212 | |
| 213 | # logical_to_physical_map: maps logical experts to physical experts |
| 214 | # Shape: (num_logical_experts, max_slots) |
| 215 | # For testing, use simple 1:1 mapping with single slot per expert |
| 216 | logical_to_physical_map = torch.arange( |
| 217 | global_num_experts, dtype=torch.int64, device="cuda" |
| 218 | ).unsqueeze(-1) |
| 219 | |
| 220 | # logical_replica_count: number of replicas per logical expert |
| 221 | # Shape: (num_logical_experts,) |
| 222 | # For testing, each logical expert has exactly 1 replica |
| 223 | logical_replica_count = torch.ones( |
| 224 | global_num_experts, dtype=torch.int64, device="cuda" |
| 225 | ) |
| 226 | should_record_tensor = torch.ones((), dtype=torch.bool, device="cuda") |
| 227 | num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32, device="cuda")] |
| 228 | |
| 229 | return EplbLayerState( |
| 230 | expert_load_view=expert_load_view, |
| 231 | logical_to_physical_map=logical_to_physical_map, |
| 232 | logical_replica_count=logical_replica_count, |
| 233 | should_record_tensor=should_record_tensor, |
| 234 | num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, |
| 235 | ) |
| 236 | |
| 237 | |
| 238 | def make_test_data( |
| 239 | m: int, k: int, num_experts: int |
| 240 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 241 | hidden_states = torch.randn((m, k), device="cuda") / 10 |
| 242 | logits = torch.randn((m, num_experts), device="cuda") |
| 243 | return hidden_states, logits |
| 244 | |
| 245 | |
| 246 | def make_e_score_correction_bias( |
| 247 | e_score_correction_bias_val: float, |
| 248 | num_experts: int, |
| 249 | ) -> torch.Tensor: |
| 250 | # return torch.randn(num_experts, device="cuda") * e_score_correction_bias_val |
| 251 | return torch.full( |
| 252 | (num_experts,), e_score_correction_bias_val, device="cuda", dtype=torch.float32 |
| 253 | ) |
| 254 | |
| 255 | |
| 256 | def assert_routing_results_close( |
| 257 | topk_weights: torch.Tensor, |
| 258 | topk_ids: torch.Tensor, |
| 259 | baseline_weights: torch.Tensor, |
| 260 | baseline_ids: torch.Tensor, |
| 261 | rtol: float = 1e-3, |
| 262 | atol: float = 1e-3, |
| 263 | ): |
| 264 | """Compare routing results, sorting by expert ID first to handle non-deterministic |
| 265 | ordering from sorted=False in topk. |
| 266 | """ |
| 267 | # Sort both results by expert IDs for consistent comparison |
| 268 | sorted_indices_actual = torch.argsort(topk_ids, dim=-1) |
| 269 | sorted_indices_baseline = torch.argsort(baseline_ids.to(topk_ids.dtype), dim=-1) |
| 270 | |
| 271 | # Gather the sorted values |
| 272 | topk_ids_sorted = torch.gather(topk_ids, 1, sorted_indices_actual) |
| 273 | topk_weights_sorted = torch.gather(topk_weights, 1, sorted_indices_actual) |
| 274 | baseline_ids_sorted = torch.gather( |
| 275 | baseline_ids.to(topk_ids.dtype), 1, sorted_indices_baseline |
| 276 | ) |
| 277 | baseline_weights_sorted = torch.gather(baseline_weights, 1, sorted_indices_baseline) |
| 278 | |
| 279 | # Compare |
| 280 | torch.testing.assert_close(topk_ids_sorted, baseline_ids_sorted) |
| 281 | torch.testing.assert_close( |
| 282 | topk_weights_sorted, baseline_weights_sorted, rtol=rtol, atol=atol |
| 283 | ) |
| 284 | |
| 285 | |
| 286 | def assert_aiter_routing_valid( |
| 287 | topk_weights: torch.Tensor, |
| 288 | topk_ids: torch.Tensor, |
| 289 | top_k: int, |
| 290 | num_experts: int, |
| 291 | renormalize: bool, |
| 292 | routed_scaling_factor: float = 1.0, |
| 293 | ): |
| 294 | """Validate AITER routing outputs are structurally correct. |
| 295 | |
| 296 | AITER grouped_topk is a fundamentally different implementation from |
| 297 | the Python baseline (different group selection, scoring internals), |
| 298 | so numerical comparison is not meaningful. Instead we verify the |
| 299 | outputs satisfy the routing contract: correct shapes, valid expert |
| 300 | IDs, non-negative weights, and proper normalization.""" |
| 301 | n_tokens = topk_weights.shape[0] |
| 302 | |
| 303 | # Shape |
| 304 | assert topk_weights.shape == (n_tokens, top_k), ( |
| 305 | f"weights shape {topk_weights.shape} != ({n_tokens}, {top_k})" |
| 306 | ) |
| 307 | assert topk_ids.shape == (n_tokens, top_k), ( |
| 308 | f"ids shape {topk_ids.shape} != ({n_tokens}, {top_k})" |
| 309 | ) |
| 310 | |
| 311 | # Expert IDs in valid range |
| 312 | assert (topk_ids >= 0).all() and (topk_ids < num_experts).all(), ( |
| 313 | f"expert IDs out of range [0, {num_experts}): " |
| 314 | f"min={topk_ids.min().item()}, max={topk_ids.max().item()}" |
| 315 | ) |
| 316 | |
| 317 | # No duplicate expert IDs per token |
| 318 | for i in range(n_tokens): |
| 319 | ids = topk_ids[i] |
| 320 | assert ids.unique().numel() == top_k, ( |
| 321 | f"token {i}: duplicate expert IDs {ids.tolist()}" |
| 322 | ) |
| 323 | |
| 324 | # Weights are non-negative |
| 325 | assert (topk_weights >= 0).all(), "negative routing weights" |
| 326 | |
| 327 | # If renormalized, weights should sum to ~scaling_factor per token |
| 328 | # (renormalization to 1.0 happens before scaling) |
| 329 | if renormalize: |
| 330 | expected_sum = routed_scaling_factor |
| 331 | sums = topk_weights.sum(dim=-1) |
| 332 | torch.testing.assert_close( |
| 333 | sums, |
| 334 | torch.full_like(sums, expected_sum), |
| 335 | atol=1e-3, |
| 336 | rtol=1e-3, |
| 337 | ) |
| 338 | |
| 339 | |
| 340 | def baseline_fused_topk( |
| 341 | router_logits: torch.Tensor, top_k: int, renormalize: bool |
| 342 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 343 | """Baseline for standard fused top-k routing. |
| 344 | |
| 345 | Algorithm: |
| 346 | 1. Apply softmax to router logits |
| 347 | 2. Select top-k experts |
| 348 | 3. Optionally renormalize the weights |
| 349 | """ |
| 350 | scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32) |
| 351 | # Use sorted=False to match vllm implementation (vllm_is_batch_invariant |
| 352 | # defaults to False) |
| 353 | topk_weights, topk_ids = torch.topk(scores, top_k, dim=-1, sorted=False) |
| 354 | |
| 355 | if renormalize: |
| 356 | topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) |
| 357 | |
| 358 | return topk_weights.to(torch.float32), topk_ids.to(torch.int32) |
| 359 | |
| 360 | |
| 361 | def baseline_fused_topk_bias( |
| 362 | router_logits: torch.Tensor, |
| 363 | top_k: int, |
| 364 | renormalize: bool, |
| 365 | e_score_correction_bias: torch.Tensor, |
| 366 | routed_scaling_factor: float, |
| 367 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 368 | """Baseline for fused top-k with bias correction. |
| 369 | |
| 370 | Algorithm: |
| 371 | 1. Apply softmax to router logits |
| 372 | 2. Add bias to scores for expert selection |
| 373 | 3. Select top-k experts using biased scores |
| 374 | 4. Get weights from original (unbiased) scores |
| 375 | 5. Apply routed scaling factor |
| 376 | 6. Optionally renormalize the weights |
| 377 | """ |
| 378 | # Apply softmax to get scores |
| 379 | scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32) |
| 380 | |
| 381 | # Add bias for expert selection |
| 382 | scores_for_choice = scores + e_score_correction_bias.unsqueeze(0) |
| 383 | |
| 384 | # Select top-k using biased scores (sorted=False to match implementation) |
| 385 | topk_ids = torch.topk(scores_for_choice, k=top_k, dim=-1, sorted=False)[1] |
| 386 | |
| 387 | # Get weights from original scores (not biased) |
| 388 | topk_weights = scores.gather(1, topk_ids) |
| 389 | |
| 390 | # Renormalize if needed (BEFORE applying scaling factor) |
| 391 | if renormalize: |
| 392 | topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) |
| 393 | |
| 394 | # Apply scaling factor (AFTER renormalization, if applicable) |
| 395 | if routed_scaling_factor != 1.0: |
| 396 | topk_weights *= routed_scaling_factor |
| 397 | |
| 398 | return topk_weights.to(torch.float32), topk_ids.to(torch.int32) |
| 399 | |
| 400 | |
| 401 | def baseline_grouped_topk( |
| 402 | router_logits: torch.Tensor, |
| 403 | top_k: int, |
| 404 | num_expert_group: int, |
| 405 | topk_group: int, |
| 406 | scoring_func: str, |
| 407 | renormalize: bool, |
| 408 | e_score_correction_bias: torch.Tensor | None, |
| 409 | routed_scaling_factor: float, |
| 410 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 411 | """Baseline for grouped top-k routing (e.g., DeepSeek). |
| 412 | |
| 413 | Algorithm: |
| 414 | 1. Apply scoring function (softmax or sigmoid) |
| 415 | 2. Optionally add bias |
| 416 | 3. Select top-k groups based on max scores within each group |
| 417 | 4. Mask scores to only include selected groups |
| 418 | 5. Select top-k experts from masked scores |
| 419 | 6. Apply scaling factor |
| 420 | 7. Optionally renormalize |
| 421 | """ |
| 422 | num_token = router_logits.shape[0] |
| 423 | |
| 424 | # Apply scoring function |
| 425 | if scoring_func == "softmax": |
| 426 | scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32) |
| 427 | elif scoring_func == "sigmoid": |
| 428 | scores = torch.sigmoid(router_logits.float()) |
| 429 | else: |
| 430 | raise ValueError(f"Unsupported scoring function: {scoring_func}") |
| 431 | |
| 432 | # Handle bias correction |
| 433 | if e_score_correction_bias is not None: |
| 434 | original_scores = scores |
| 435 | scores = scores + e_score_correction_bias.unsqueeze(0) |
| 436 | # For bias case, use sum of top-2 scores in each group |
| 437 | group_scores = ( |
| 438 | scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1) |
| 439 | ) |
| 440 | else: |
| 441 | # Use max score in each group |
| 442 | group_scores = scores.view(num_token, num_expert_group, -1).max(dim=-1).values |
| 443 | |
| 444 | # Select top-k groups |
| 445 | group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1] |
| 446 | |
| 447 | # Create mask for selected groups |
| 448 | group_mask = torch.zeros_like(group_scores) |
| 449 | group_mask.scatter_(1, group_idx, 1) |
| 450 | |
| 451 | # Expand mask to all experts |
| 452 | score_mask = ( |
| 453 | group_mask.unsqueeze(-1) |
| 454 | .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) |
| 455 | .reshape(num_token, -1) |
| 456 | ) |
| 457 | |
| 458 | # Mask scores (set non-selected to -inf) |
| 459 | tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) |
| 460 | |
| 461 | # Select top-k experts |
| 462 | if e_score_correction_bias is not None: |
| 463 | topk_ids = torch.topk(tmp_scores, k=top_k, dim=-1, sorted=False)[1] |
| 464 | topk_weights = original_scores.gather(1, topk_ids) |
| 465 | else: |
| 466 | topk_weights, topk_ids = torch.topk(tmp_scores, k=top_k, dim=-1, sorted=False) |
| 467 | |
| 468 | # Renormalize if needed |
| 469 | if renormalize: |
| 470 | topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) |
| 471 | |
| 472 | # Apply scaling factor |
| 473 | if routed_scaling_factor != 1.0: |
| 474 | topk_weights *= routed_scaling_factor |
| 475 | |
| 476 | return topk_weights.to(torch.float32), topk_ids.to(torch.int32) |
| 477 | |
| 478 | |
| 479 | def baseline_custom_llama4( |
| 480 | router_logits: torch.Tensor, top_k: int |
| 481 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 482 | """Baseline for Llama4 custom routing. |
| 483 | |
| 484 | Algorithm: |
| 485 | 1. Select top-k expert indices (without softmax) |
| 486 | 2. Apply sigmoid to the selected scores |
| 487 | """ |
| 488 | router_scores, router_indices = torch.topk(router_logits, top_k, dim=-1) |
| 489 | router_scores = torch.sigmoid(router_scores.float()) |
| 490 | return router_scores.to(torch.float32), router_indices.to(torch.int32) |
| 491 | |
| 492 | |
| 493 | @pytest.mark.parametrize("m,k", MK_S) |
| 494 | @pytest.mark.parametrize("top_k", TOP_KS) |
| 495 | @pytest.mark.parametrize("global_num_experts", NUM_EXPERTS) |
| 496 | @pytest.mark.parametrize("renormalize", [False, True]) |
| 497 | @pytest.mark.parametrize("enable_eplb", [False, True]) |
| 498 | def test_fused_topk( |
| 499 | m: int, |
| 500 | k: int, |
| 501 | top_k: int, |
| 502 | global_num_experts: int, |
| 503 | renormalize: bool, |
| 504 | enable_eplb: bool, |
| 505 | ): |
| 506 | if top_k > global_num_experts: |
| 507 | pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})") |
| 508 | |
| 509 | eplb_state = setup_eplb_state(enable_eplb, global_num_experts) |
| 510 | router = create_fused_moe_router( |
| 511 | top_k=top_k, |
| 512 | global_num_experts=global_num_experts, |
| 513 | renormalize=renormalize, |
| 514 | eplb_state=eplb_state, |
| 515 | ) |
| 516 | |
| 517 | hidden_states, router_logits = make_test_data(m, k, global_num_experts) |
| 518 | |
| 519 | # Get router output |
| 520 | topk_weights, topk_ids = router.select_experts(hidden_states, router_logits) |
| 521 | |
| 522 | # Compute baseline |
| 523 | baseline_weights, baseline_ids = baseline_fused_topk( |
| 524 | router_logits, top_k, renormalize |
| 525 | ) |
| 526 | |
| 527 | # Compare results |
| 528 | assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids) |
| 529 | |
| 530 | |
| 531 | @pytest.mark.parametrize("m,k", MK_S) |
| 532 | @pytest.mark.parametrize("top_k", TOP_KS) |
| 533 | @pytest.mark.parametrize("global_num_experts", NUM_EXPERTS) |
| 534 | @pytest.mark.parametrize("renormalize", [False, True]) |
| 535 | @pytest.mark.parametrize("enable_eplb", [False, True]) |
| 536 | @pytest.mark.parametrize("e_score_correction_bias_val", [0.9]) |
| 537 | @pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.1]) |
| 538 | def test_fused_topk_bias( |
| 539 | m: int, |
| 540 | k: int, |
| 541 | top_k: int, |
| 542 | global_num_experts: int, |
| 543 | renormalize: bool, |
| 544 | enable_eplb: bool, |
| 545 | e_score_correction_bias_val: float, |
| 546 | routed_scaling_factor: float, |
| 547 | ): |
| 548 | if top_k > global_num_experts: |
| 549 | pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})") |
| 550 | |
| 551 | eplb_state = setup_eplb_state(enable_eplb, global_num_experts) |
| 552 | |
| 553 | e_score_correction_bias = make_e_score_correction_bias( |
| 554 | e_score_correction_bias_val, |
| 555 | global_num_experts, |
| 556 | ) |
| 557 | |
| 558 | router = create_fused_moe_router( |
| 559 | e_score_correction_bias=e_score_correction_bias, |
| 560 | routed_scaling_factor=routed_scaling_factor, |
| 561 | top_k=top_k, |
| 562 | global_num_experts=global_num_experts, |
| 563 | renormalize=renormalize, |
| 564 | eplb_state=eplb_state, |
| 565 | ) |
| 566 | |
| 567 | hidden_states, router_logits = make_test_data(m, k, global_num_experts) |
| 568 | |
| 569 | # Get router output |
| 570 | topk_weights, topk_ids = router.select_experts(hidden_states, router_logits) |
| 571 | |
| 572 | # Compute baseline |
| 573 | baseline_weights, baseline_ids = baseline_fused_topk_bias( |
| 574 | router_logits, |
| 575 | top_k, |
| 576 | renormalize, |
| 577 | e_score_correction_bias, |
| 578 | routed_scaling_factor, |
| 579 | ) |
| 580 | |
| 581 | # Compare results |
| 582 | assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids) |
| 583 | |
| 584 | |
| 585 | @pytest.mark.parametrize("m,k", MK_S) |
| 586 | @pytest.mark.parametrize("top_k", TOP_KS) |
| 587 | @pytest.mark.parametrize( |
| 588 | "global_num_experts,num_expert_group,topk_group", |
| 589 | [ |
| 590 | (64, 8, 4), # 8 groups of 8 experts, select 4 groups |
| 591 | (32, 4, 2), # 4 groups of 8 experts, select 2 groups |
| 592 | ], |
| 593 | ) |
| 594 | @pytest.mark.parametrize("renormalize", [False, True]) |
| 595 | @pytest.mark.parametrize("enable_eplb", [False, True]) |
| 596 | @pytest.mark.parametrize("e_score_correction_bias_val", [0.9]) |
| 597 | @pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.1]) |
| 598 | @pytest.mark.parametrize("scoring_func", ["sigmoid", "softmax"]) |
| 599 | def test_grouped_topk( |
| 600 | m: int, |
| 601 | k: int, |
| 602 | top_k: int, |
| 603 | global_num_experts: int, |
| 604 | renormalize: bool, |
| 605 | enable_eplb: bool, |
| 606 | num_expert_group: int, |
| 607 | topk_group: int, |
| 608 | scoring_func: str, |
| 609 | e_score_correction_bias_val: float, |
| 610 | routed_scaling_factor: float, |
| 611 | ): |
| 612 | if top_k > global_num_experts: |
| 613 | pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})") |
| 614 | |
| 615 | eplb_state = setup_eplb_state(enable_eplb, global_num_experts) |
| 616 | |
| 617 | e_score_correction_bias = make_e_score_correction_bias( |
| 618 | e_score_correction_bias_val, |
| 619 | global_num_experts, |
| 620 | ) |
| 621 | |
| 622 | router = create_fused_moe_router( |
| 623 | use_grouped_topk=True, |
| 624 | num_expert_group=num_expert_group, |
| 625 | topk_group=topk_group, |
| 626 | scoring_func=scoring_func, |
| 627 | e_score_correction_bias=e_score_correction_bias, |
| 628 | routed_scaling_factor=routed_scaling_factor, |
| 629 | top_k=top_k, |
| 630 | global_num_experts=global_num_experts, |
| 631 | renormalize=renormalize, |
| 632 | eplb_state=eplb_state, |
| 633 | ) |
| 634 | |
| 635 | hidden_states, router_logits = make_test_data(m, k, global_num_experts) |
| 636 | |
| 637 | # Compute baseline (pure Python implementation) |
| 638 | baseline_weights, baseline_ids = baseline_grouped_topk( |
| 639 | router_logits, |
| 640 | top_k, |
| 641 | num_expert_group, |
| 642 | topk_group, |
| 643 | scoring_func, |
| 644 | renormalize, |
| 645 | e_score_correction_bias, |
| 646 | routed_scaling_factor, |
| 647 | ) |
| 648 | |
| 649 | # Test 1: Python/Triton path against baseline (exact match) |
| 650 | with patch( |
| 651 | "vllm.model_executor.layers.fused_moe.router.grouped_topk_router.rocm_aiter_ops.is_fused_moe_enabled", |
| 652 | return_value=False, |
| 653 | ): |
| 654 | py_weights, py_ids = router.select_experts(hidden_states, router_logits) |
| 655 | assert_routing_results_close(py_weights, py_ids, baseline_weights, baseline_ids) |
| 656 | |
| 657 | # Test 2: AITER path — verify outputs are structurally valid. |
| 658 | # AITER grouped_topk is a different implementation so we can't |
| 659 | # compare numerically against the Python baseline. |
| 660 | if _is_aiter_capable(): |
| 661 | # Force-enable AITER for gfx942/gfx950 regardless of env var, |
| 662 | # so CI always exercises this path on capable hardware. |
| 663 | with patch.object(rocm_aiter_ops, "_AITER_ENABLED", True): |
| 664 | aiter_weights, aiter_ids = router.select_experts( |
| 665 | hidden_states, router_logits |
| 666 | ) |
| 667 | assert_aiter_routing_valid( |
| 668 | aiter_weights, |
| 669 | aiter_ids, |
| 670 | top_k, |
| 671 | global_num_experts, |
| 672 | renormalize, |
| 673 | routed_scaling_factor, |
| 674 | ) |
| 675 | |
| 676 | |
| 677 | @pytest.mark.parametrize("m,k", MK_S) |
| 678 | @pytest.mark.parametrize("top_k", TOP_KS) |
| 679 | @pytest.mark.parametrize("global_num_experts", NUM_EXPERTS) |
| 680 | @pytest.mark.parametrize("renormalize", [False, True]) |
| 681 | @pytest.mark.parametrize("enable_eplb", [False, True]) |
| 682 | @pytest.mark.parametrize("custom_routing_function", [Llama4MoE.custom_routing_function]) |
| 683 | def test_custom( |
| 684 | m: int, |
| 685 | k: int, |
| 686 | top_k: int, |
| 687 | global_num_experts: int, |
| 688 | renormalize: bool, |
| 689 | enable_eplb: bool, |
| 690 | custom_routing_function: Callable, |
| 691 | ): |
| 692 | if top_k > global_num_experts: |
| 693 | pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})") |
| 694 | |
| 695 | eplb_state = setup_eplb_state(enable_eplb, global_num_experts) |
| 696 | |
| 697 | router = create_fused_moe_router( |
| 698 | top_k=top_k, |
| 699 | global_num_experts=global_num_experts, |
| 700 | custom_routing_function=custom_routing_function, |
| 701 | renormalize=renormalize, |
| 702 | eplb_state=eplb_state, |
| 703 | ) |
| 704 | |
| 705 | hidden_states, router_logits = make_test_data(m, k, global_num_experts) |
| 706 | |
| 707 | # Get router output |
| 708 | topk_weights, topk_ids = router.select_experts(hidden_states, router_logits) |
| 709 | |
| 710 | # Compute baseline (Llama4 uses sigmoid) |
| 711 | baseline_weights, baseline_ids = baseline_custom_llama4(router_logits, top_k) |
| 712 | |
| 713 | # Compare results |
| 714 | assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids) |
| 715 | |
| 716 | |
| 717 | # TODO: is other test sufficient? |
| 718 | # # See tests/test_routing_simulatator.py |
| 719 | # @pytest.mark.parametrize("m,k", MK_S) |
| 720 | # @pytest.mark.parametrize("top_k", TOP_KS) |
| 721 | # @pytest.mark.parametrize("global_num_experts", NUM_EXPERTS) |
| 722 | # @pytest.mark.parametrize("renormalize", [False, True]) |
| 723 | # @pytest.mark.parametrize("enable_eplb", [False, True]) |
| 724 | # @pytest.mark.parameterize("strategy", ["uniform_random", "normal_routing"]) |
| 725 | # def test_simulated( |
| 726 | # m: int, |
| 727 | # k: int, |
| 728 | # top_k: int, |
| 729 | # global_num_experts: int, |
| 730 | # renormalize: bool, |
| 731 | # enable_eplb: bool, |
| 732 | # strategy: str, |
| 733 | # monkeypatch, |
| 734 | # ): |
| 735 | # eplb_state = setup_eplb_state(enable_eplb) |
| 736 | |
| 737 | # monkeypatch.setenv("VLLM_MOE_ROUTING_SIMULATION_STRATEGY", strategy) |
| 738 | # router = create_fused_moe_router( |
| 739 | # top_k=top_k, |
| 740 | # global_num_experts=global_num_experts, |
| 741 | # eplb_state=eplb_state, |
| 742 | # ) |
| 743 | |
| 744 | # hidden_states, router_logits = make_test_data(m, k, global_num_experts) |
| 745 | # topk_weights, topk_ids = router.select_experts(hidden_states, router_logits) |
| 746 | |
| 747 | |
| 748 | # --------------------------------------------------------------------------- |
| 749 | # Tests for eplb_map_to_physical_and_record |
| 750 | # --------------------------------------------------------------------------- |
| 751 | |
| 752 | |
| 753 | @pytest.mark.parametrize("record_enabled", [True, False]) |
| 754 | @pytest.mark.parametrize( |
| 755 | "l2p_map, replica_count, num_physical, topk_ids, expected_out, expected_load", |
| 756 | [ |
| 757 | pytest.param( |
| 758 | # logical i → physical i |
| 759 | [[0], [1], [2], [3]], |
| 760 | [1, 1, 1, 1], |
| 761 | 4, |
| 762 | [[0, 1], [2, 3], [0, 2]], |
| 763 | [[0, 1], [2, 3], [0, 2]], |
| 764 | [2, 1, 2, 1], |
| 765 | id="identity", |
| 766 | ), |
| 767 | pytest.param( |
| 768 | # logical 0→3, 1→0, 2→1, 3→2 |
| 769 | [[3], [0], [1], [2]], |
| 770 | [1, 1, 1, 1], |
| 771 | 4, |
| 772 | [[0, 1], [2, 3], [0, 2]], |
| 773 | [[3, 0], [1, 2], [3, 1]], |
| 774 | [1, 2, 1, 2], |
| 775 | id="shuffled", |
| 776 | ), |
| 777 | pytest.param( |
| 778 | # logical 0→5, 1→2, 2→7, 3→0 in a larger physical space |
| 779 | [[5], [2], [7], [0]], |
| 780 | [1, 1, 1, 1], |
| 781 | 8, |
| 782 | [[0, 1], [2, 3]], |
| 783 | [[5, 2], [7, 0]], |
| 784 | [1, 0, 1, 0, 0, 1, 0, 1], |
| 785 | id="sparse", |
| 786 | ), |
| 787 | ], |
| 788 | ) |
| 789 | def test_eplb_map_no_redundancy( |
| 790 | record_enabled, |
| 791 | l2p_map, |
| 792 | replica_count, |
| 793 | num_physical, |
| 794 | topk_ids, |
| 795 | expected_out, |
| 796 | expected_load, |
| 797 | ): |
| 798 | l2p = torch.tensor(l2p_map, dtype=torch.int64, device="cuda") |
| 799 | rc = torch.tensor(replica_count, dtype=torch.int64, device="cuda") |
| 800 | load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") |
| 801 | rec = torch.tensor(record_enabled, dtype=torch.bool, device="cuda") |
| 802 | ids = torch.tensor(topk_ids, dtype=torch.int32, device="cuda") |
| 803 | |
| 804 | out = eplb_map_to_physical_and_record( |
| 805 | topk_ids=ids, |
| 806 | expert_load_view=load, |
| 807 | logical_to_physical_map=l2p, |
| 808 | logical_replica_count=rc, |
| 809 | record_enabled=rec, |
| 810 | ) |
| 811 | |
| 812 | exp_out = torch.tensor(expected_out, dtype=out.dtype, device="cuda") |
| 813 | torch.testing.assert_close(out, exp_out) |
| 814 | |
| 815 | if record_enabled: |
| 816 | exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda") |
| 817 | torch.testing.assert_close(load, exp_load) |
| 818 | else: |
| 819 | assert load.sum().item() == 0 |
| 820 | |
| 821 | |
| 822 | @pytest.mark.parametrize("top_k,R", [(2, 2), (4, 2), (8, 4), (8, 8)]) |
| 823 | def test_eplb_map_hot_expert_replica_balance(top_k, R): |
| 824 | """Hot logical expert with R replicas must be balanced across replicas |
| 825 | even when ``top_k`` is a multiple of ``R``. In that regime every top-k |
| 826 | offset for the hot expert lands on a multiple of ``top_k`` in the flat |
| 827 | ``topk_ids`` view, so per-replica assignment must not collapse onto a |
| 828 | single replica. |
| 829 | """ |
| 830 | num_tokens = 8192 |
| 831 | num_logical = 16 |
| 832 | num_physical = R + (num_logical - 1) |
| 833 | |
| 834 | l2p = torch.full((num_logical, R), -1, dtype=torch.int64, device="cuda") |
| 835 | l2p[0] = torch.arange(R, dtype=torch.int64, device="cuda") |
| 836 | for i in range(1, num_logical): |
| 837 | l2p[i, 0] = R + i - 1 |
| 838 | rc = torch.tensor([R] + [1] * (num_logical - 1), dtype=torch.int64, device="cuda") |
| 839 | |
| 840 | torch.manual_seed(0) |
| 841 | topk_ids = torch.randint( |
| 842 | 1, |
| 843 | num_logical, |
| 844 | (num_tokens, top_k), |
| 845 | dtype=torch.int32, |
| 846 | device="cuda", |
| 847 | ) |
| 848 | topk_ids[:, 0] = 0 |
| 849 | |
| 850 | load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") |
| 851 | rec = torch.tensor(True, dtype=torch.bool, device="cuda") |
| 852 | |
| 853 | eplb_map_to_physical_and_record( |
| 854 | topk_ids=topk_ids, |
| 855 | expert_load_view=load, |
| 856 | logical_to_physical_map=l2p, |
| 857 | logical_replica_count=rc, |
| 858 | record_enabled=rec, |
| 859 | ) |
| 860 | |
| 861 | hot_load = load[:R].float() |
| 862 | max_mean = (hot_load.max() / hot_load.mean()).item() |
| 863 | assert max_mean < 1.15, ( |
| 864 | f"Hot expert replicas uneven: {hot_load.tolist()}, max/mean={max_mean:.3f}" |
| 865 | ) |
| 866 | |
| 867 | |
| 868 | @pytest.mark.parametrize("record_enabled", [True, False]) |
| 869 | @pytest.mark.parametrize( |
| 870 | "l2p_map, replica_count, num_physical, topk_ids, expected_out, expected_load", |
| 871 | [ |
| 872 | pytest.param( |
| 873 | # experts 0,1 have 2 replicas; 2,3 have 1 |
| 874 | [[0, 4], [1, 5], [2, -1], [3, -1]], |
| 875 | [2, 2, 1, 1], |
| 876 | 6, |
| 877 | [[0, 1], [2, 3], [0, 2]], |
| 878 | # replica = (token_idx * KNUTH) & 0xFFFFFFFF % R. |
| 879 | # token 0 hash=0x00000000: %2=0, %1=0. |
| 880 | # token 1 hash=0x9E3779B9: %2=1, %1=0. |
| 881 | # token 2 hash=0x3C6EF372: %2=0, %1=0. |
| 882 | [[0, 1], [2, 3], [0, 2]], |
| 883 | [2, 1, 2, 1, 0, 0], |
| 884 | id="partial", |
| 885 | ), |
| 886 | pytest.param( |
| 887 | # all 4 experts have 2 replicas |
| 888 | [[0, 4], [1, 5], [2, 6], [3, 7]], |
| 889 | [2, 2, 2, 2], |
| 890 | 8, |
| 891 | [[0, 1], [2, 3], [0, 2]], |
| 892 | # token 0 hash=0x00000000: %2=0. |
| 893 | # token 1 hash=0x9E3779B9: %2=1. |
| 894 | # token 2 hash=0x3C6EF372: %2=0. |
| 895 | [[0, 1], [6, 7], [0, 2]], |
| 896 | [2, 1, 1, 0, 0, 0, 1, 1], |
| 897 | id="full", |
| 898 | ), |
| 899 | pytest.param( |
| 900 | # expert 0: 4 replicas, experts 1,2: 2 replicas |
| 901 | [[0, 3, 5, 7], [1, 4, -1, -1], [2, 6, -1, -1]], |
| 902 | [4, 2, 2], |
| 903 | 8, |
| 904 | [[0, 1], [2, 0], [1, 2]], |
| 905 | # token 0 hash=0x00000000: %4=0, %2=0. |
| 906 | # token 1 hash=0x9E3779B9: %4=1, %2=1. |
| 907 | # token 2 hash=0x3C6EF372: %4=2, %2=0. |
| 908 | [[0, 1], [6, 3], [1, 2]], |
| 909 | [1, 2, 1, 1, 0, 0, 1, 0], |
| 910 | id="uneven", |
| 911 | ), |
| 912 | ], |
| 913 | ) |
| 914 | def test_eplb_map_with_redundancy( |
| 915 | record_enabled, |
| 916 | l2p_map, |
| 917 | replica_count, |
| 918 | num_physical, |
| 919 | topk_ids, |
| 920 | expected_out, |
| 921 | expected_load, |
| 922 | ): |
| 923 | l2p = torch.tensor(l2p_map, dtype=torch.int64, device="cuda") |
| 924 | rc = torch.tensor(replica_count, dtype=torch.int64, device="cuda") |
| 925 | load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") |
| 926 | rec = torch.tensor(record_enabled, dtype=torch.bool, device="cuda") |
| 927 | ids = torch.tensor(topk_ids, dtype=torch.int32, device="cuda") |
| 928 | |
| 929 | out = eplb_map_to_physical_and_record( |
| 930 | topk_ids=ids, |
| 931 | expert_load_view=load, |
| 932 | logical_to_physical_map=l2p, |
| 933 | logical_replica_count=rc, |
| 934 | record_enabled=rec, |
| 935 | ) |
| 936 | |
| 937 | exp_out = torch.tensor(expected_out, dtype=out.dtype, device="cuda") |
| 938 | torch.testing.assert_close(out, exp_out) |
| 939 | |
| 940 | if record_enabled: |
| 941 | exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda") |
| 942 | torch.testing.assert_close(load, exp_load) |
| 943 | else: |
| 944 | assert load.sum().item() == 0 |
| 945 | |
| 946 | |
| 947 | @pytest.mark.parametrize( |
| 948 | "l2p_map, replica_count, num_physical, topk_ids, " |
| 949 | "num_unpadded, expected_out, expected_load", |
| 950 | [ |
| 951 | pytest.param( |
| 952 | [[0], [1], [2], [3]], |
| 953 | [1, 1, 1, 1], |
| 954 | 4, |
| 955 | [[0, 1], [2, 3], [0, 2], [1, 3]], |
| 956 | 2, |
| 957 | [[0, 1], [2, 3], [0, 2], [1, 3]], |
| 958 | # only rows 0,1 counted: expert 0→1, 1→1, 2→1, 3→1 |
| 959 | [1, 1, 1, 1], |
| 960 | id="half_padded", |
| 961 | ), |
| 962 | pytest.param( |
| 963 | # record everything (None = no padding info) |
| 964 | [[0], [1], [2], [3]], |
| 965 | [1, 1, 1, 1], |
| 966 | 4, |
| 967 | [[0, 1], [2, 3], [0, 2], [1, 3]], |
| 968 | None, |
| 969 | [[0, 1], [2, 3], [0, 2], [1, 3]], |
| 970 | [2, 2, 2, 2], |
| 971 | id="no_padding_info", |
| 972 | ), |
| 973 | ], |
| 974 | ) |
| 975 | def test_eplb_map_num_unpadded_tokens( |
| 976 | l2p_map, |
| 977 | replica_count, |
| 978 | num_physical, |
| 979 | topk_ids, |
| 980 | num_unpadded, |
| 981 | expected_out, |
| 982 | expected_load, |
| 983 | ): |
| 984 | l2p = torch.tensor(l2p_map, dtype=torch.int64, device="cuda") |
| 985 | rc = torch.tensor(replica_count, dtype=torch.int64, device="cuda") |
| 986 | load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") |
| 987 | rec = torch.tensor(True, dtype=torch.bool, device="cuda") |
| 988 | ids = torch.tensor(topk_ids, dtype=torch.int32, device="cuda") |
| 989 | num_unpadded_t = ( |
| 990 | torch.tensor(num_unpadded, dtype=torch.int32, device="cuda") |
| 991 | if num_unpadded is not None |
| 992 | else None |
| 993 | ) |
| 994 | |
| 995 | out = eplb_map_to_physical_and_record( |
| 996 | topk_ids=ids, |
| 997 | expert_load_view=load, |
| 998 | logical_to_physical_map=l2p, |
| 999 | logical_replica_count=rc, |
| 1000 | record_enabled=rec, |
| 1001 | num_unpadded_tokens=num_unpadded_t, |
| 1002 | ) |
| 1003 | |
| 1004 | exp_out = torch.tensor(expected_out, dtype=out.dtype, device="cuda") |
| 1005 | torch.testing.assert_close(out, exp_out) |
| 1006 | |
| 1007 | exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda") |
| 1008 | torch.testing.assert_close(load, exp_load) |