[MoE Refactor] Separate Router into OO Classes (#30623)

Signed-off-by: Bill Nell <bnell@redhat.com>
2026-01-18 11:40:49 -05:00
parent 2f03035a61
commit 327a02d8db
45 changed files with 1750 additions and 684 deletions
--- a/tests/kernels/moe/test_routing.py
+++ b/tests/kernels/moe/test_routing.py
@@ -0,0 +1,499 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
+
+import pytest
+import torch
+
+from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
+from vllm.model_executor.layers.fused_moe.router.router_factory import (
+    create_fused_moe_router,
+)
+from vllm.model_executor.models.llama4 import Llama4MoE
+
+# Test parameters
+MK_S = [(32, 256), (64, 512)]
+TOP_KS = [2, 4, 6]
+NUM_EXPERTS = [8, 16, 64]
+
+
+def setup_eplb_state(enable_eplb: bool, global_num_experts: int) -> EplbLayerState:
+    if not enable_eplb:
+        return EplbLayerState()
+
+    # Initialize EPLB state with proper tensors for testing
+    # For testing purposes, we use a simple 1:1 mapping (no redundant experts)
+    # expert_load_view: tracks load on each expert (shape: num_experts)
+    expert_load_view = torch.zeros(global_num_experts, dtype=torch.int32, device="cuda")
+
+    # logical_to_physical_map: maps logical experts to physical experts
+    # Shape: (num_logical_experts, max_slots)
+    # For testing, use simple 1:1 mapping with single slot per expert
+    logical_to_physical_map = torch.arange(
+        global_num_experts, dtype=torch.int64, device="cuda"
+    ).unsqueeze(-1)
+
+    # logical_replica_count: number of replicas per logical expert
+    # Shape: (num_logical_experts,)
+    # For testing, each logical expert has exactly 1 replica
+    logical_replica_count = torch.ones(
+        global_num_experts, dtype=torch.int64, device="cuda"
+    )
+
+    return EplbLayerState(
+        expert_load_view=expert_load_view,
+        logical_to_physical_map=logical_to_physical_map,
+        logical_replica_count=logical_replica_count,
+    )
+
+
+def make_test_data(
+    m: int, k: int, num_experts: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    hidden_states = torch.randn((m, k), device="cuda") / 10
+    logits = torch.randn((m, num_experts), device="cuda")
+    return hidden_states, logits
+
+
+def make_e_score_correction_bias(
+    e_score_correction_bias_val: float,
+    num_experts: int,
+) -> torch.Tensor:
+    # return torch.randn(num_experts, device="cuda") * e_score_correction_bias_val
+    return torch.full(
+        (num_experts,), e_score_correction_bias_val, device="cuda", dtype=torch.float32
+    )
+
+
+def assert_routing_results_close(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    baseline_weights: torch.Tensor,
+    baseline_ids: torch.Tensor,
+    rtol: float = 1e-3,
+    atol: float = 1e-3,
+):
+    """
+    Compare routing results, sorting by expert ID first to handle non-deterministic
+    ordering from sorted=False in topk.
+    """
+    # Sort both results by expert IDs for consistent comparison
+    sorted_indices_actual = torch.argsort(topk_ids, dim=-1)
+    sorted_indices_baseline = torch.argsort(baseline_ids.to(topk_ids.dtype), dim=-1)
+
+    # Gather the sorted values
+    topk_ids_sorted = torch.gather(topk_ids, 1, sorted_indices_actual)
+    topk_weights_sorted = torch.gather(topk_weights, 1, sorted_indices_actual)
+    baseline_ids_sorted = torch.gather(
+        baseline_ids.to(topk_ids.dtype), 1, sorted_indices_baseline
+    )
+    baseline_weights_sorted = torch.gather(baseline_weights, 1, sorted_indices_baseline)
+
+    # Compare
+    torch.testing.assert_close(topk_ids_sorted, baseline_ids_sorted)
+    torch.testing.assert_close(
+        topk_weights_sorted, baseline_weights_sorted, rtol=rtol, atol=atol
+    )
+
+
+def baseline_fused_topk(
+    router_logits: torch.Tensor, top_k: int, renormalize: bool
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Baseline for standard fused top-k routing.
+
+    Algorithm:
+    1. Apply softmax to router logits
+    2. Select top-k experts
+    3. Optionally renormalize the weights
+    """
+    scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
+    # Use sorted=False to match vllm implementation (vllm_is_batch_invariant
+    # defaults to False)
+    topk_weights, topk_ids = torch.topk(scores, top_k, dim=-1, sorted=False)
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+
+
+def baseline_fused_topk_bias(
+    router_logits: torch.Tensor,
+    top_k: int,
+    renormalize: bool,
+    e_score_correction_bias: torch.Tensor,
+    routed_scaling_factor: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Baseline for fused top-k with bias correction.
+
+    Algorithm:
+    1. Apply softmax to router logits
+    2. Add bias to scores for expert selection
+    3. Select top-k experts using biased scores
+    4. Get weights from original (unbiased) scores
+    5. Apply routed scaling factor
+    6. Optionally renormalize the weights
+    """
+    # Apply softmax to get scores
+    scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
+
+    # Add bias for expert selection
+    scores_for_choice = scores + e_score_correction_bias.unsqueeze(0)
+
+    # Select top-k using biased scores (sorted=False to match implementation)
+    topk_ids = torch.topk(scores_for_choice, k=top_k, dim=-1, sorted=False)[1]
+
+    # Get weights from original scores (not biased)
+    topk_weights = scores.gather(1, topk_ids)
+
+    # Renormalize if needed (BEFORE applying scaling factor)
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    # Apply scaling factor (AFTER renormalization, if applicable)
+    if routed_scaling_factor != 1.0:
+        topk_weights *= routed_scaling_factor
+
+    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+
+
+def baseline_grouped_topk(
+    router_logits: torch.Tensor,
+    top_k: int,
+    num_expert_group: int,
+    topk_group: int,
+    scoring_func: str,
+    renormalize: bool,
+    e_score_correction_bias: torch.Tensor | None,
+    routed_scaling_factor: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Baseline for grouped top-k routing (e.g., DeepSeek).
+
+    Algorithm:
+    1. Apply scoring function (softmax or sigmoid)
+    2. Optionally add bias
+    3. Select top-k groups based on max scores within each group
+    4. Mask scores to only include selected groups
+    5. Select top-k experts from masked scores
+    6. Apply scaling factor
+    7. Optionally renormalize
+    """
+    num_token = router_logits.shape[0]
+
+    # Apply scoring function
+    if scoring_func == "softmax":
+        scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
+    elif scoring_func == "sigmoid":
+        scores = torch.sigmoid(router_logits.float())
+    else:
+        raise ValueError(f"Unsupported scoring function: {scoring_func}")
+
+    # Handle bias correction
+    if e_score_correction_bias is not None:
+        original_scores = scores
+        scores = scores + e_score_correction_bias.unsqueeze(0)
+        # For bias case, use sum of top-2 scores in each group
+        group_scores = (
+            scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+        )
+    else:
+        # Use max score in each group
+        group_scores = scores.view(num_token, num_expert_group, -1).max(dim=-1).values
+
+    # Select top-k groups
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1]
+
+    # Create mask for selected groups
+    group_mask = torch.zeros_like(group_scores)
+    group_mask.scatter_(1, group_idx, 1)
+
+    # Expand mask to all experts
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
+        .reshape(num_token, -1)
+    )
+
+    # Mask scores (set non-selected to -inf)
+    tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf"))
+
+    # Select top-k experts
+    if e_score_correction_bias is not None:
+        topk_ids = torch.topk(tmp_scores, k=top_k, dim=-1, sorted=False)[1]
+        topk_weights = original_scores.gather(1, topk_ids)
+    else:
+        topk_weights, topk_ids = torch.topk(tmp_scores, k=top_k, dim=-1, sorted=False)
+
+    # Renormalize if needed
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    # Apply scaling factor
+    if routed_scaling_factor != 1.0:
+        topk_weights *= routed_scaling_factor
+
+    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+
+
+def baseline_custom_llama4(
+    router_logits: torch.Tensor, top_k: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Baseline for Llama4 custom routing.
+
+    Algorithm:
+    1. Select top-k expert indices (without softmax)
+    2. Apply sigmoid to the selected scores
+    """
+    router_scores, router_indices = torch.topk(router_logits, top_k, dim=-1)
+    router_scores = torch.sigmoid(router_scores.float())
+    return router_scores.to(torch.float32), router_indices.to(torch.int32)
+
+
+@pytest.mark.parametrize("m,k", MK_S)
+@pytest.mark.parametrize("top_k", TOP_KS)
+@pytest.mark.parametrize("global_num_experts", NUM_EXPERTS)
+@pytest.mark.parametrize("renormalize", [False, True])
+@pytest.mark.parametrize("enable_eplb", [False, True])
+def test_fused_topk(
+    m: int,
+    k: int,
+    top_k: int,
+    global_num_experts: int,
+    renormalize: bool,
+    enable_eplb: bool,
+):
+    if top_k > global_num_experts:
+        pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})")
+
+    eplb_state = setup_eplb_state(enable_eplb, global_num_experts)
+    router = create_fused_moe_router(
+        top_k=top_k,
+        global_num_experts=global_num_experts,
+        renormalize=renormalize,
+        enable_eplb=enable_eplb,
+        eplb_state=eplb_state,
+    )
+
+    hidden_states, router_logits = make_test_data(m, k, global_num_experts)
+
+    # Get router output
+    topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
+
+    # Compute baseline
+    baseline_weights, baseline_ids = baseline_fused_topk(
+        router_logits, top_k, renormalize
+    )
+
+    # Compare results
+    assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
+
+
+@pytest.mark.parametrize("m,k", MK_S)
+@pytest.mark.parametrize("top_k", TOP_KS)
+@pytest.mark.parametrize("global_num_experts", NUM_EXPERTS)
+@pytest.mark.parametrize("renormalize", [False, True])
+@pytest.mark.parametrize("enable_eplb", [False, True])
+@pytest.mark.parametrize("e_score_correction_bias_val", [0.9])
+@pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.1])
+def test_fused_topk_bias(
+    m: int,
+    k: int,
+    top_k: int,
+    global_num_experts: int,
+    renormalize: bool,
+    enable_eplb: bool,
+    e_score_correction_bias_val: float,
+    routed_scaling_factor: float,
+):
+    if top_k > global_num_experts:
+        pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})")
+
+    eplb_state = setup_eplb_state(enable_eplb, global_num_experts)
+
+    e_score_correction_bias = make_e_score_correction_bias(
+        e_score_correction_bias_val,
+        global_num_experts,
+    )
+
+    router = create_fused_moe_router(
+        e_score_correction_bias=e_score_correction_bias,
+        routed_scaling_factor=routed_scaling_factor,
+        top_k=top_k,
+        global_num_experts=global_num_experts,
+        renormalize=renormalize,
+        enable_eplb=enable_eplb,
+        eplb_state=eplb_state,
+    )
+
+    hidden_states, router_logits = make_test_data(m, k, global_num_experts)
+
+    # Get router output
+    topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
+
+    # Compute baseline
+    baseline_weights, baseline_ids = baseline_fused_topk_bias(
+        router_logits,
+        top_k,
+        renormalize,
+        e_score_correction_bias,
+        routed_scaling_factor,
+    )
+
+    # Compare results
+    assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
+
+
+@pytest.mark.parametrize("m,k", MK_S)
+@pytest.mark.parametrize("top_k", TOP_KS)
+@pytest.mark.parametrize(
+    "global_num_experts,num_expert_group,topk_group",
+    [
+        (64, 8, 4),  # 8 groups of 8 experts, select 4 groups
+        (32, 4, 2),  # 4 groups of 8 experts, select 2 groups
+    ],
+)
+@pytest.mark.parametrize("renormalize", [False, True])
+@pytest.mark.parametrize("enable_eplb", [False, True])
+@pytest.mark.parametrize("e_score_correction_bias_val", [0.9])
+@pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.1])
+@pytest.mark.parametrize("scoring_func", ["sigmoid", "softmax"])
+def test_grouped_topk(
+    m: int,
+    k: int,
+    top_k: int,
+    global_num_experts: int,
+    renormalize: bool,
+    enable_eplb: bool,
+    num_expert_group: int,
+    topk_group: int,
+    scoring_func: str,
+    e_score_correction_bias_val: float,
+    routed_scaling_factor: float,
+):
+    if top_k > global_num_experts:
+        pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})")
+
+    eplb_state = setup_eplb_state(enable_eplb, global_num_experts)
+
+    e_score_correction_bias = make_e_score_correction_bias(
+        e_score_correction_bias_val,
+        global_num_experts,
+    )
+
+    routing_method_type = None
+    if scoring_func == "llama4":
+        routing_method_type = RoutingMethodType.Llama4
+        scoring_func = "sigmoid"
+
+    router = create_fused_moe_router(
+        use_grouped_topk=True,
+        num_expert_group=num_expert_group,
+        topk_group=topk_group,
+        scoring_func=scoring_func,
+        routing_method_type=routing_method_type,
+        e_score_correction_bias=e_score_correction_bias,
+        routed_scaling_factor=routed_scaling_factor,
+        top_k=top_k,
+        global_num_experts=global_num_experts,
+        renormalize=renormalize,
+        enable_eplb=enable_eplb,
+        eplb_state=eplb_state,
+    )
+
+    hidden_states, router_logits = make_test_data(m, k, global_num_experts)
+
+    # Get router output
+    topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
+
+    # Compute baseline
+    baseline_weights, baseline_ids = baseline_grouped_topk(
+        router_logits,
+        top_k,
+        num_expert_group,
+        topk_group,
+        scoring_func,
+        renormalize,
+        e_score_correction_bias,
+        routed_scaling_factor,
+    )
+
+    # Compare results
+    assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
+
+
+@pytest.mark.parametrize("m,k", MK_S)
+@pytest.mark.parametrize("top_k", TOP_KS)
+@pytest.mark.parametrize("global_num_experts", NUM_EXPERTS)
+@pytest.mark.parametrize("renormalize", [False, True])
+@pytest.mark.parametrize("enable_eplb", [False, True])
+@pytest.mark.parametrize("custom_routing_function", [Llama4MoE.custom_routing_function])
+def test_custom(
+    m: int,
+    k: int,
+    top_k: int,
+    global_num_experts: int,
+    renormalize: bool,
+    enable_eplb: bool,
+    custom_routing_function: Callable,
+):
+    if top_k > global_num_experts:
+        pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})")
+
+    eplb_state = setup_eplb_state(enable_eplb, global_num_experts)
+
+    router = create_fused_moe_router(
+        top_k=top_k,
+        global_num_experts=global_num_experts,
+        custom_routing_function=custom_routing_function,
+        renormalize=renormalize,
+        enable_eplb=enable_eplb,
+        eplb_state=eplb_state,
+    )
+
+    hidden_states, router_logits = make_test_data(m, k, global_num_experts)
+
+    # Get router output
+    topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
+
+    # Compute baseline (Llama4 uses sigmoid)
+    baseline_weights, baseline_ids = baseline_custom_llama4(router_logits, top_k)
+
+    # Compare results
+    assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
+
+
+# TODO: is other test sufficient?
+# # See tests/test_routing_simulatator.py
+# @pytest.mark.parametrize("m,k", MK_S)
+# @pytest.mark.parametrize("top_k", TOP_KS)
+# @pytest.mark.parametrize("global_num_experts", NUM_EXPERTS)
+# @pytest.mark.parametrize("renormalize", [False, True])
+# @pytest.mark.parametrize("enable_eplb", [False, True])
+# @pytest.mark.parameterize("strategy", ["uniform_random", "normal_routing"])
+# def test_simulated(
+#     m: int,
+#     k: int,
+#     top_k: int,
+#     global_num_experts: int,
+#     renormalize: bool,
+#     enable_eplb: bool,
+#     strategy: str,
+#     monkeypatch,
+# ):
+#     eplb_state = setup_eplb_state(enable_eplb)
+
+#     monkeypatch.setenv("VLLM_MOE_ROUTING_SIMULATION_STRATEGY", strategy)
+#     router = create_fused_moe_router(
+#         top_k=top_k,
+#         global_num_experts=global_num_experts,
+#         enable_eplb=enable_eplb,
+#         eplb_state=eplb_state,
+#     )
+
+#     hidden_states, router_logits = make_test_data(m, k, global_num_experts)
+#     topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)