Repository navigation
Order NaN as the largest value in the CPU TopK kernel - #33543
Merged
mirounga merged 1 commit intoOct 29, 2026
Merged
Conversation
With a NaN in the input, the CPU comparators were not a strict weak ordering, so the selected elements were arbitrary: TopK of [1, NaN, 3, 2] with k = 2 returned [1, NaN]. NaN now ranks above every other value, as it does in the CUDA kernel, and ties keep the lower index first. The k = 1 scan keeps its plain comparison and only rescans rows that contain a NaN. Co-Authored-By: Claude Opus 5.5 <[email protected]>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
mirounga
approved these changes
Oct 29, 2026
This was referenced Nov 9, 2026
1 task done
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.
Description
the CPU
TopKcomparators use>/<and==on the values , so with a NaN in the input they are not a strict weak ordering : a NaN is neither greater , smaller nor equal to anything . the result depends on where the NaN sits and which path runs , andstd::nth_element/std::sortare handed a comparator they do not support .[1, NaN, 3, 2][1, NaN][NaN, 3][1, NaN, 3, 2, NaN, -1, 5, 0]5at 6NaNat 1NaN now ranks above every other value and equal values , NaN included , keep the lower index first . that is what the CUDA kernel already does ,
HybridNaNUsesStablePackedOrderingchecks it , so CPU and CUDA agree . integer types compile to the same comparisons as before .the
k == 1scan keeps its plain comparison , which the compiler vectorizes , and folds a NaN check into the same loop . only a row that holds a NaN is scanned again with the NaN aware comparison . a standalone copy of the loops on 4M floats without NaN , clang -O2 on arm64 :k == 1scanputting the NaN check straight into the
k == 1comparison instead made that loop about 6x slower , which is why it has the second pass .TopKOperator.NaNIsLargestcovers the scan ( k = 1 ) , the heap ( k = 3 ) andnth_element( k = 6 and 8 ) for largest and smallest , float and double , opset 11 and 24 , on the CPU EP . it fails before the change .onnxruntime_provider_test --gtest_filter='*TopK*:*BeamSearch*:*GreedySearch*:*Sampling*', 100 tests , passes on macOS arm64 , andlintrunneris clean .Motivation and Context
found by running single node models through the ONNX reference evaluator and onnxruntime 1.31 side by side . TopK of
[1, NaN, 3, 2]with k = 2 gave[1, NaN]here , 1 being the smallest value in the input . the spec does not define a NaN order , so this follows the CUDA EP . no existing issue that I could find .