Skip to content

Add support for stable sorted outputs to DeviceBatchedTopK #11326

Open

Description

@pauleonix

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    • Status
      Todo

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions