test_routing.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
from unittest.mock import patch

import pytest
import torch

from vllm._aiter_ops import rocm_aiter_ops
from vllm.distributed.eplb.eplb_state import EplbLayerState
from vllm.model_executor.layers.fused_moe.router.router_factory import (
    create_fused_moe_router,
)
from vllm.model_executor.models.llama4 import Llama4MoE
from vllm.platforms import current_platform


def _is_aiter_capable() -> bool:
    """Check if the platform supports AITER (gfx942/gfx950)."""
    if not current_platform.is_rocm():
        return False
    try:
        from vllm.platforms.rocm import _ON_MI3XX

        return _ON_MI3XX
    except ImportError:
        return False


# 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 assert_aiter_routing_valid(
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    top_k: int,
    num_experts: int,
    renormalize: bool,
    routed_scaling_factor: float = 1.0,
):
    """Validate AITER routing outputs are structurally correct.

    AITER grouped_topk is a fundamentally different implementation from
    the Python baseline (different group selection, scoring internals),
    so numerical comparison is not meaningful. Instead we verify the
    outputs satisfy the routing contract: correct shapes, valid expert
    IDs, non-negative weights, and proper normalization."""
    n_tokens = topk_weights.shape[0]

    # Shape
    assert topk_weights.shape == (n_tokens, top_k), (
        f"weights shape {topk_weights.shape} != ({n_tokens}, {top_k})"
    )
    assert topk_ids.shape == (n_tokens, top_k), (
        f"ids shape {topk_ids.shape} != ({n_tokens}, {top_k})"
    )

    # Expert IDs in valid range
    assert (topk_ids >= 0).all() and (topk_ids < num_experts).all(), (
        f"expert IDs out of range [0, {num_experts}): "
        f"min={topk_ids.min().item()}, max={topk_ids.max().item()}"
    )

    # No duplicate expert IDs per token
    for i in range(n_tokens):
        ids = topk_ids[i]
        assert ids.unique().numel() == top_k, (
            f"token {i}: duplicate expert IDs {ids.tolist()}"
        )

    # Weights are non-negative
    assert (topk_weights >= 0).all(), "negative routing weights"

    # If renormalized, weights should sum to ~scaling_factor per token
    # (renormalization to 1.0 happens before scaling)
    if renormalize:
        expected_sum = routed_scaling_factor
        sums = topk_weights.sum(dim=-1)
        torch.testing.assert_close(
            sums,
            torch.full_like(sums, expected_sum),
            atol=1e-3,
            rtol=1e-3,
        )


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,
    )

    router = create_fused_moe_router(
        use_grouped_topk=True,
        num_expert_group=num_expert_group,
        topk_group=topk_group,
        scoring_func=scoring_func,
        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)

    # Compute baseline (pure Python implementation)
    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,
    )

    # Test 1: Python/Triton path against baseline (exact match)
    with patch(
        "vllm.model_executor.layers.fused_moe.router.grouped_topk_router.rocm_aiter_ops.is_fused_moe_enabled",
        return_value=False,
    ):
        py_weights, py_ids = router.select_experts(hidden_states, router_logits)
    assert_routing_results_close(py_weights, py_ids, baseline_weights, baseline_ids)

    # Test 2: AITER path — verify outputs are structurally valid.
    # AITER grouped_topk is a different implementation so we can't
    # compare numerically against the Python baseline.
    if _is_aiter_capable():
        # Force-enable AITER for gfx942/gfx950 regardless of env var,
        # so CI always exercises this path on capable hardware.
        with patch.object(rocm_aiter_ops, "_AITER_ENABLED", True):
            aiter_weights, aiter_ids = router.select_experts(
                hidden_states, router_logits
            )
        assert_aiter_routing_valid(
            aiter_weights,
            aiter_ids,
            top_k,
            global_num_experts,
            renormalize,
            routed_scaling_factor,
        )


@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)