Repository navigation
tests: fix CI flakes of TOPK_MOE caused by near ties - #31426
Open
ServeurpersoCom wants to merge 1 commit into
Open
ServeurpersoCom wants to merge 1 commit into
ServeurpersoCom wants to merge 1 commit into
Conversation
With uniform random logits, two experts sometimes land within one f32 ulp of each other at a routing boundary, mostly with sqrt_softplus gating whose output range is narrow. A backend whose gating differs by a few ulps from the CPU then picks the other expert and the test fails. Draw each row of logits again until the group selection and the final top-k both have a relative margin of 1e-4, computed on the host in double precision.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
TOPK_MOE in test-backend-ops fails now and then on master, almost always with sqrt_softplus gating, recently on the WebGPU Apple and Vulkan Intel runners. One expert index differs from the CPU, and the next run passes.
The test checks that a backend picks exactly the same experts as the CPU. That question only has an answer when the choice is clear. With uniform random logits, two experts sometimes end up within one f32 ulp of each other right at the boundary: the last expert kept by the top-k, or the last group kept in grouped routing. Either choice is then equally valid, and the result depends on how each backend rounds exp, log or sqrt in its last bit. sqrt_softplus makes this more frequent because it squeezes the logits into a narrow range.
This PR keeps the random data but draws a row again when one of its routing decisions is that close to a tie (margin below 1e-4 relative, checked on the host in double precision). The kernels run the same code on the same kind of data, the comparison stays strict, and the test no longer asks backends to agree on a coin flip.
Additional information
The fragile cases come from the sqrt_softplus gating added to this test in #25896, and the larger expert counts added since (#25267, #31288) make near ties more likely.
The flake also shows up on a rootless container with an RTX PRO 6000 on Vulkan (5 failed runs out of 148). A temporary probe that recomputed the routing on the host showed that every captured failure had a margin of 1e-7 or less, one of them an exact tie. With this PR, 1400 runs of the sqrt_softplus cases pass without a failure, where about 11 were expected before, and TOPK_MOE passes on Vulkan, WebGPU and CUDA.
Requirements