From dde54168c6258c2e4c7c9a22ae8c72ac024552b2 Mon Sep 17 00:00:00 2001 From: David Canar Date: Sat, 5 Sep 2026 18:35:54 -0600 Subject: [PATCH] [Bugfix][MoE] Deterministic expert selection in the grouped_topk Python fallback The Python `grouped_topk` fallback selects experts with `torch.topk(..., sorted=envs.VLLM_BATCH_INVARIANT)`, i.e. `sorted=False` by default. `sorted=False` licenses the backend to return the k results in any order, and above 256 columns `torch.topk` takes a multi-pass path whose order is not merely unsorted but differs between calls on identical input: 20 identical `torch.topk(x, k=8, sorted=False)` calls on a 64xE tensor give 1 distinct result for E <= 256 and 20 distinct for E >= 257 (gfx1151, ROCm 10.0, torch 2.11; the same split holds for any k >= 4). Two consequences for MoE routing: 1. The routing weights are permuted to match the returned order, so the order the expert outputs are summed in changes run to run and the MoE output is not reproducible. On a real 740x288 GLM-5.3 router score tensor this affected 740/740 rows on every call. 2. On exact fp32 ties at the k-boundary the selected set itself flips. `sorted=True` does not fix that, and `torch.use_deterministic_algorithms(True)` neither raises, warns, nor changes the behaviour. Select with a stable descending sort instead, giving a total order of value descending, then expert index ascending. This is not a new convention: it is what the fused CUDA kernel already implements (`moeTopKFuncs.cuh` packs `65535 - idx` into the comparison key, and the multi-group path prefers the lower index on equal values), what `test_grouped_topk_single_group_stable_ties` already asserts, and what `_single_group_reference` in that same test file already computes with `argsort(..., descending=True, stable=True)`. The Python fallback was the only piece that disagreed -- and it is the path ROCm takes, as well as the path CUDA takes whenever `e_score_correction_bias is None` or the shape falls outside the fused kernel's tier table. Because it needs more than 256 routed experts to appear at all, this has been easy to miss: DeepSeek-V3/R1 has exactly 256 and sits on the safe side of the boundary, while GLM-5.3's 288 is over it. End to end on GLM-5.3-Flash (TP=2, 2x gfx1151), 5 byte-identical greedy requests per prompt length, distinct completions before -> after: 244 tok 5 -> 1, 614 tok 5 -> 1, 1196 tok 3 -> 1, 1421 tok 5 -> 1, 1797 tok 5 -> 1. With a hash tap after every decoder layer, all 45 layers become bit-identical across repeated forwards on both TP ranks. Cost is negative: 71.6us vs 97.8us for the selection on the real 740x288 tensor. The new tests are deliberately not CUDA-gated, unlike the rest of that file, because the pure-PyTorch fallback is the code under test. They use 288 experts because the same tests written at 256 pass against unfixed code: 16 of the 18 new cases fail before this change and all 18 pass after. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: David Canar --- tests/kernels/moe/test_grouped_topk.py | 335 ++++++++++++++++++ .../fused_moe/router/grouped_topk_router.py | 26 +- 2 files changed, 352 insertions(+), 9 deletions(-) diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py index 1ab8550d64e8..96608f3b81ca 100644 --- a/tests/kernels/moe/test_grouped_topk.py +++ b/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:]) diff --git a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py index 1d4e4a8b5e22..b74eb7846279 100644 --- a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py +++ b/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)