[Bugfix][MoE] Deterministic expert selection in the grouped_topk Python fallback

tests/kernels/moe/test_grouped_topk.py

@@ -18,6 +18,7 @@
 from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
     GroupedTopk,
     fused_grouped_topk,
+    grouped_topk,
 )
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
@@ -367,3 +368,337 @@ def test_grouped_topk_single_group_nonfinite_scores(
 
     torch.testing.assert_close(actual_ids, expected_ids)
     torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# Determinism of the pure-PyTorch fallback (`grouped_topk`)
+#
+# Everything above exercises the fused CUDA kernel and is therefore skipped on
+# non-CUDA platforms. The tests below deliberately are NOT skipped: the code
+# under test is the Python fallback, which is the live path on ROCm (whenever
+# AITER MoE is off) and is also reached on CUDA whenever the fused kernel is
+# not taken -- `VLLM_USE_FUSED_MOE_GROUPED_TOPK=0`, `e_score_correction_bias
+# is None`, or a shape outside the kernel's tier table.
+#
+# The fallback used to select with `torch.topk(..., sorted=False)`, and
+# `sorted=False` licenses the backend to return the k results in any order.
+# Above 256 columns `torch.topk` takes a multi-pass path whose order is not
+# merely unsorted but differs between calls on identical input. Measured on
+# gfx1151 / torch 2.11+rocm10.0, 20 identical `torch.topk(x, k=8,
+# sorted=False)` calls on a 64xE tensor gave 1 distinct result for E <= 256
+# and 20 distinct for E >= 257, with the same split for any k >= 4.
+#
+# Hence _DET_NUM_EXPERTS = 288 below: at E <= 256 -- which includes
+# DeepSeek-V3/R1's 256, sitting exactly on the safe side of the boundary --
+# the stock implementation happens to return sorted, stable output and these
+# tests pass against unfixed code.
+# ---------------------------------------------------------------------------
+
+# The fallback is device-agnostic; run on whatever the runner has. (On ROCm
+# builds torch.cuda.is_available() is True and "cuda" is the right device
+# string -- HIP presents itself through the CUDA API.)
+_DET_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+_DET_NUM_REPEAT = 20
+_DET_NUM_EXPERTS = 288
+_DET_TOPK = 8
+_DET_NUM_TOKENS = 64
+
+
+def _run_python_grouped_topk(
+    logits: torch.Tensor,
+    bias: torch.Tensor | None,
+    topk: int,
+    *,
+    num_expert_group: int = 1,
+    topk_group: int = 1,
+    scoring_func: str = "sigmoid",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Call the Python fallback exactly as the router does when the fused
+    kernel is not taken."""
+    return grouped_topk(
+        hidden_states=torch.empty(
+            (logits.shape[0], 0), dtype=logits.dtype, device=logits.device
+        ),
+        gating_output=logits,
+        topk=topk,
+        renormalize=False,
+        num_expert_group=num_expert_group,
+        topk_group=topk_group,
+        scoring_func=scoring_func,
+        e_score_correction_bias=bias,
+    )
+
+
+def _det_scores(logits: torch.Tensor, scoring_func: str) -> torch.Tensor:
+    if scoring_func == "sigmoid":
+        return logits.sigmoid()
+    return torch.softmax(logits, dim=-1)
+
+
+def _det_biased_scores(
+    logits: torch.Tensor, bias: torch.Tensor | None, scoring_func: str
+) -> torch.Tensor:
+    scores = _det_scores(logits, scoring_func)
+    if bias is not None:
+        scores = scores + bias.unsqueeze(0)
+    return scores
+
+
+def _force_python_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Force the Python path on every platform (including CUDA), and make it
+    # explicit that determinism must not require VLLM_BATCH_INVARIANT, which
+    # is documented as NVIDIA SM90+ only and so cannot be enabled on every
+    # stack that reaches this code.
+    monkeypatch.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
+
+
+def _det_random_inputs(
+    bias_is_none: bool, seed: int = 0
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    gen = torch.Generator(device="cpu").manual_seed(seed)
+    logits = torch.randn(
+        _DET_NUM_TOKENS, _DET_NUM_EXPERTS, generator=gen, dtype=torch.float32
+    ).to(_DET_DEVICE)
+    bias = (
+        None
+        if bias_is_none
+        else torch.randn(_DET_NUM_EXPERTS, generator=gen, dtype=torch.float32).to(
+            _DET_DEVICE
+        )
+    )
+    return logits, bias
+
+
+def _assert_all_identical(
+    results: list[tuple[torch.Tensor, torch.Tensor]],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    first_w, first_i = results[0]
+    for n, (w, i) in enumerate(results[1:], start=1):
+        assert torch.equal(i, first_i), (
+            f"expert ids differ between call 0 and call {n}: "
+            f"{(i != first_i).sum().item()} of {i.numel()} positions"
+        )
+        # Bitwise, not merely value-wise: this also catches a -0.0/0.0 flip
+        # and any reordering among equal weights.
+        assert w.view(torch.int32).equal(first_w.view(torch.int32)), (
+            f"routing weights differ between call 0 and call {n}"
+        )
+    return first_w, first_i
+
+
+def _make_tie_logits(
+    num_experts: int, k: int, tie_lo: int, tie_hi: int
+) -> torch.Tensor:
+    """One row with a deliberate, bitwise-exact tie at the k-boundary.
+
+    Experts ``0 .. k-2`` get well-separated descending logits, experts
+    ``tie_lo`` (= k-1) and ``tie_hi`` get the *same* logit value, and every
+    other expert sits strictly below the pair. Because the tied pair has
+    identical input bits and the scoring activation is elementwise, the two
+    biased scores are bitwise-equal on every backend, eager or compiled --
+    the tie is constructed, not a rounding accident. The pair occupies
+    positions k-1 and k of the value-descending order, i.e. the last selected
+    slot and the first dropped one.
+    """
+    assert tie_lo == k - 1 and tie_hi > tie_lo and tie_hi < num_experts
+    logits = torch.zeros(num_experts, dtype=torch.float32)
+    logits[: k - 1] = torch.linspace(8.0, 4.0, k - 1)
+    logits[tie_lo] = 3.0
+    logits[tie_hi] = 3.0
+    rest = torch.ones(num_experts, dtype=torch.bool)
+    rest[: k - 1] = False
+    rest[tie_lo] = False
+    rest[tie_hi] = False
+    logits[rest] = torch.linspace(2.9, 0.1, int(rest.sum()))
+    return logits
+
+
+@pytest.mark.parametrize("scoring_func", ["sigmoid", "softmax"])
+@pytest.mark.parametrize("bias_is_none", [False, True])
+@pytest.mark.parametrize(
+    ("num_expert_group", "topk_group"), [(1, 1), (8, 4)], ids=["1grp", "8grp"]
+)
+def test_grouped_topk_repeat_determinism(
+    monkeypatch: pytest.MonkeyPatch,
+    scoring_func: str,
+    bias_is_none: bool,
+    num_expert_group: int,
+    topk_group: int,
+):
+    """Identical inputs must give bitwise-identical routing.
+
+    No tie is needed: with ``sorted=False`` the fallback returned the same k
+    experts in a *different order* on every call once E > 256, which permutes
+    the routing weights and so changes the order the expert outputs are summed
+    in downstream.
+    """
+    _force_python_fallback(monkeypatch)
+    logits, bias = _det_random_inputs(bias_is_none)
+
+    # A fresh clone per call so the input buffer address varies, as it does in
+    # a real forward.
+    results = [
+        _run_python_grouped_topk(
+            logits.clone(),
+            bias,
+            _DET_TOPK,
+            num_expert_group=num_expert_group,
+            topk_group=topk_group,
+            scoring_func=scoring_func,
+        )
+        for _ in range(_DET_NUM_REPEAT)
+    ]
+    _assert_all_identical(results)
+
+
+@pytest.mark.parametrize("scoring_func", ["sigmoid", "softmax"])
+@pytest.mark.parametrize("bias_is_none", [False, True])
+def test_grouped_topk_returns_value_descending_order(
+    monkeypatch: pytest.MonkeyPatch, scoring_func: str, bias_is_none: bool
+):
+    """The selected experts must come back in descending score order.
+
+    This is the order the fused kernel already produces -- ``moeTopKFuncs.cuh``
+    packs ``65535 - idx`` into the comparison key and the multi-group path uses
+    ``WarpSelect<..., is_stable=true>`` -- and the order
+    ``_single_group_reference`` above already assumes, via a stable descending
+    ``argsort``. It is a single-call assertion, so it cannot be flaky.
+    """
+    _force_python_fallback(monkeypatch)
+    logits, bias = _det_random_inputs(bias_is_none)
+    biased = _det_biased_scores(logits, bias, scoring_func)
+
+    _, topk_ids = _run_python_grouped_topk(
+        logits, bias, _DET_TOPK, scoring_func=scoring_func
+    )
+
+    selected = biased.gather(1, topk_ids.to(torch.long))
+    bad = (selected[:, :-1] < selected[:, 1:]).any(dim=1)
+    assert not bool(bad.any()), (
+        f"{int(bad.sum())} of {bad.numel()} rows are not in descending score "
+        f"order; first offending row {int(bad.nonzero()[0])}: "
+        f"{selected[int(bad.nonzero()[0])].tolist()}"
+    )
+
+    # Same experts as the reference: on tie-free input this order is exactly
+    # what topk(..., sorted=True) already returns, i.e. pinning the order does
+    # not change the selection.
+    ref_ids = biased.topk(_DET_TOPK, dim=-1, sorted=True)[1].to(torch.int32)
+    torch.testing.assert_close(topk_ids, ref_ids)
+
+
+@pytest.mark.parametrize("scoring_func", ["sigmoid", "softmax"])
+@pytest.mark.parametrize("bias_is_none", [False, True])
+def test_grouped_topk_tie_broken_by_lower_expert_index(
+    monkeypatch: pytest.MonkeyPatch, scoring_func: str, bias_is_none: bool
+):
+    """An exact tie at the k-boundary is resolved by the lower expert index.
+
+    This is the same contract as ``test_grouped_topk_single_group_stable_ties``
+    asserts for the fused kernel, but for the fallback. Such ties do occur in
+    practice: a live 740x288 router score tensor from a GLM-5.3-Flash prefill
+    had k-boundary ties on 2 of 740 rows, between experts whose logits *and*
+    biases both differ but whose fp32 sums round to the same value.
+    """
+    _force_python_fallback(monkeypatch)
+
+    tie_lo, tie_hi = _DET_TOPK - 1, 200
+    logits = _make_tie_logits(_DET_NUM_EXPERTS, _DET_TOPK, tie_lo, tie_hi)[None].to(
+        _DET_DEVICE
+    )
+    bias = None if bias_is_none else torch.zeros(_DET_NUM_EXPERTS, device=_DET_DEVICE)
+
+    biased = _det_biased_scores(logits, bias, scoring_func)
+    # Self-validate the construction before asserting anything about the op.
+    tie_val = biased[0, tie_lo]
+    assert biased[0, tie_lo].item() == biased[0, tie_hi].item()
+    assert int((biased[0] > tie_val).sum()) == _DET_TOPK - 1
+    assert int((biased[0] == tie_val).sum()) == 2
+
+    results = [
+        _run_python_grouped_topk(
+            logits.clone(), bias, _DET_TOPK, scoring_func=scoring_func
+        )
+        for _ in range(_DET_NUM_REPEAT)
+    ]
+    first_w, first_i = _assert_all_identical(results)
+
+    assert int(first_i[0, _DET_TOPK - 1]) == tie_lo
+    assert tie_hi not in first_i[0].tolist()
+
+    # The full selection is the value-descending one: experts 0..k-2 by
+    # construction, then the tie winner tie_lo = k-1.
+    expected_ids = torch.arange(_DET_TOPK, dtype=torch.int32, device=_DET_DEVICE)[None]
+    torch.testing.assert_close(first_i, expected_ids)
+
+    # Weights are the unbiased scores of the selected experts. A small
+    # tolerance vs the eager reference: the compiled elementwise activation
+    # may differ from eager by an ULP on some backends.
+    expected_w = _det_scores(logits, scoring_func).gather(1, first_i.to(torch.long))
+    torch.testing.assert_close(first_w, expected_w, atol=2e-6, rtol=0)
+
+
+@pytest.mark.parametrize("scoring_func", ["sigmoid", "softmax"])
+def test_grouped_topk_group_tie_broken_by_lower_group_index(
+    monkeypatch: pytest.MonkeyPatch, scoring_func: str
+):
+    """A bitwise-exact tie at the *group* boundary is broken by ascending
+    group index, covering the group-selection site.
+
+    This one stays small on purpose: the group top-k is over
+    ``num_expert_group`` values (8 for real DeepSeek/GLM configs), always far
+    below topk's 256-column threshold, so a genuine tie is the only way to
+    make the group-selection site observable.
+    """
+    _force_python_fallback(monkeypatch)
+
+    # 4 groups of 8 experts. Group 0 is clearly best (its top-2 dominate),
+    # groups 1 and 2 are bitwise-identical -- a deliberate exact tie for the
+    # second of the two selected groups -- and group 3 is clearly worst. The
+    # non-top logits inside each group are kept low so that the top-4
+    # individuals of the union {group 0, group 1} are exactly experts
+    # 0, 1, 8, 9; had group 2 won the tie they would be 0, 1, 16, 17.
+    g0 = [8.0, 7.0, 1.9, 1.8, 1.7, 1.6, 1.5, 1.4]
+    g12 = [4.0, 3.5, 1.3, 1.2, 1.1, 1.0, 0.95, 0.9]
+    g3 = [0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1]
+    logits = torch.tensor(g0 + g12 + g12 + g3, dtype=torch.float32)[None].to(
+        _DET_DEVICE
+    )
+    num_experts = logits.shape[1]
+    bias = torch.zeros(num_experts, device=_DET_DEVICE)
+
+    # The group scores (top-2 sum within each group) of groups 1 and 2 must be
+    # bitwise-equal: identical inputs, elementwise activation, values-only
+    # reduction. (Ties inside that top-2-by-value reduction are between equal
+    # values, so the sum is unaffected -- which is why that reduction needs no
+    # tie-break of its own.)
+    def group_scores(row: torch.Tensor) -> torch.Tensor:
+        return row.view(1, 4, -1).topk(2, dim=-1)[0].sum(dim=-1)
+
+    gs = group_scores(_det_biased_scores(logits, bias, scoring_func))
+    assert gs[0, 1].item() == gs[0, 2].item()
+    assert gs[0, 0] > gs[0, 1] and gs[0, 3] < gs[0, 1]
+
+    results = [
+        _run_python_grouped_topk(
+            logits.clone(),
+            bias,
+            4,
+            num_expert_group=4,
+            topk_group=2,
+            scoring_func=scoring_func,
+        )
+        for _ in range(_DET_NUM_REPEAT)
+    ]
+    _, first_i = _assert_all_identical(results)
+
+    # Group 1 (the lower index) wins the tie, so the experts are drawn from
+    # groups 0 and 1 only.
+    expected_ids = torch.tensor([[0, 1, 8, 9]], dtype=torch.int32, device=_DET_DEVICE)
+    torch.testing.assert_close(first_i, expected_ids)
+
+    order = _det_biased_scores(logits, bias, scoring_func).gather(
+        1, first_i.to(torch.long)
+    )
+    assert torch.all(order[:, :-1] >= order[:, 1:])

vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py

@@ -71,6 +71,18 @@ def fused_grouped_topk(
     return topk_values, topk_indices
 
 
+def _deterministic_topk_indices(x: torch.Tensor, k: int) -> torch.Tensor:
+    """Top-k indices under a total order: value descending, index ascending.
+
+    `torch.topk` guarantees no tie-break, and on some backends is
+    nondeterministic on exact ties; with `sorted=False` it additionally
+    permutes the returned order run-to-run. Both make expert selection
+    irreproducible. A stable descending sort gives the same ordering the fused
+    kernel uses (value desc, then lower index first).
+    """
+    return x.sort(dim=-1, descending=True, stable=True).indices[..., :k]
+
+
 # This is used by the Deepseek-V2 and Deepseek-V3 model
 @torch.compile(
     dynamic=True,
@@ -130,11 +142,8 @@ def grouped_topk(
             scores.view(num_token, num_expert_group, -1).max(dim=-1).values
         )  # [n, n_group]
 
-    # For batch invariance, use sorted=True to ensure deterministic expert selection
-    use_sorted = envs.VLLM_BATCH_INVARIANT
-    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=use_sorted)[
-        1
-    ]  # [n, top_k_group]
+    # [n, top_k_group]
+    group_idx = _deterministic_topk_indices(group_scores, topk_group)
     group_mask = torch.zeros_like(group_scores)  # [n, n_group]
     group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
     score_mask = (
@@ -145,13 +154,12 @@ def grouped_topk(
     tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf"))  # [n, e]
 
     if e_score_correction_bias is not None:
-        topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=use_sorted)[1]
+        topk_ids = _deterministic_topk_indices(tmp_scores, topk)
         # Use original unbiased scores for the routing weights
         topk_weights = original_scores.gather(1, topk_ids)
     else:
-        topk_weights, topk_ids = torch.topk(
-            tmp_scores, k=topk, dim=-1, sorted=use_sorted
-        )
+        topk_ids = _deterministic_topk_indices(tmp_scores, topk)
+        topk_weights = tmp_scores.gather(1, topk_ids)
 
     if renormalize:
         topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)