You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Add support for stable sorted outputs to DeviceBatchedTopK #11326
DeviceBatchedTopK supports non-stable sorted output (output_ordering::sorted) since #10812 (#10811). What is still missing is output_ordering::stable_sorted: output sorted by key, with equal keys in source-index order (smaller index first). DeviceBatchedTopK currently rejects it at compile time.
Because stable_sorted is part of the default requirements (determinism::gpu_to_gpu + tie_break::prefer_smaller_index + output_ordering::stable_sorted), calling DeviceBatchedTopK without explicit requirements does not compile either, and users have to opt out of stable output explicitly.
As with sorted, the implementation has to remain compatible with CUDA Graph capture.
This issue can be closed once output_ordering::stable_sorted is supported for the same (determinism, tie_break) combinations, segment sizes, and GPU architectures that output_ordering::sorted supports, and calling DeviceBatchedTopK without explicit requirements compiles.
Activity
DeviceBatchedTopK#9260