feat: Adaptive topk algorithm selection based on input characteristics

[ClementLinCF]

Motivation

We need a selection mechanism to choose the optimal top-K algorithm.

Technical Details

Adaptive algorithm selection based on K value:

• Small K: Bitonic sort with AMD GPU-specific optimizations (DPP, med3, double buffering)
• Large K: 11-bit radix sort for consistent performance

Test Plan

python test_topk_plain.py

Test Result

This figures illustrates the performance comparison across different Top-K strategies. The dashed line represents BlockTopK (our bitonic sort-based implementation), which demonstrates competitive performance for K ≤ 128 but begins to lag behind the radix sort-based approach (radix_11bits) as K increases beyond this threshold.

Our AdaptiveTopK strategy intelligently selects the optimal algorithm based on K: leveraging BlockTopK for K ≤ 128 and switching to radix_11bits for K > 128. This adaptive selection ensures peak performance across the entire K spectrum, as evidenced by the solid line consistently matching or exceeding the best-performing algorithm at each K value.

Submission Checklist

• [ ✓] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

This PR implements adaptive Top-K algorithm selection based on input characteristics (primarily K value) to optimize performance across different workload scenarios. The implementation intelligently chooses between bitonic sort for small K values (≤128) and 11-bit radix sort for larger K values, maximizing performance across the entire spectrum.

Key Changes:

• Adaptive algorithm selection heuristic based on K value and input length
• Extended capacity support from 512 to 2048 elements
• Variable-length input support via rowStarts/rowEnds parameters

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
op_tests/test_topk_plain.py	Updated test to handle new function signatures and expanded test parameter ranges; added exception handling for Triton fallback
csrc/kernels/topk_plain_kernels.cu	Core implementation of adaptive selection logic; added support for variable-length inputs and increased MAX_CAPACITY to 2048
csrc/kernels/topk_per_row_kernels.cu	Added null-safe handling for rowStarts/rowEnds parameters and extended k parameter to filtering functions
csrc/include/topk_plain.h	Updated function signature to include topk_out tensor and variable-length support parameters
csrc/include/rocm_ops.hpp	Updated Python bindings with new parameters and default values
csrc/include/opus/opus.hpp	Fixed type inconsistency: changed __fp16 to _Float16 for med3 template specialization
aiter/ops/topk_plain.py	Updated Python wrapper signature with additional parameters for variable-length support
aiter/jit/optCompilerConfig.json	Added topk_per_row_kernels.cu to build sources

Comments suppressed due to low confidence (1)

csrc/kernels/topk_per_row_kernels.cu:1340

• Similar to the other function, the parameter k is added to filter_and_histogram_for_one_block at line 1182 but is never used in the function body (lines 1183-1340). This indicates either an incomplete implementation or an unnecessary parameter addition. Review if k should be used for any logic within this function.

                                                   IdxT k)
{
    constexpr int num_buckets = calc_num_buckets<BitsPerPass>();
    for(int i = threadIdx.x; i < num_buckets * 2; i += blockDim.x)
    {
        histogram[i] = 0;
    }
    IdxT* p_filter_cnt = &counter->filter_cnt;
    if(threadIdx.x == 0)
    {
        *p_filter_cnt = 0;
    }
    __syncthreads();

    int const start_bit = calc_start_bit<T, BitsPerPass>(pass);
    unsigned const mask = calc_mask<T, BitsPerPass>(pass);

    if(pass == 0)
    {
        T local_min = std::numeric_limits<T>::max();
        T local_max = std::numeric_limits<T>::lowest();

        auto f = [histogram, select_min, start_bit, mask, &local_min, &local_max](
                     T value, IdxT, int& acc, int& prev_bin_idx, bool is_last) {
            int bucket = calc_bucket<T, BitsPerPass>(value, start_bit, mask, select_min);
            // atomicAdd(histogram + bucket, static_cast<IdxT>(1));

            if(bucket == prev_bin_idx)
            {
                acc++;
            }
            else
            {
                if(acc > 0)
                {
                    atomicAdd(histogram + prev_bin_idx, static_cast<IdxT>(acc));
                }
                acc          = 1;
                prev_bin_idx = bucket;
            }

            if(is_last)
            {
                return;
            }

            int bucket_low =
                calc_bucket<T, BitsPerPass>(value, 0, (1 << BitsPerPass) - 1, select_min);
            atomicAdd(histogram + num_buckets + bucket_low, static_cast<IdxT>(1));

            local_min = fminf(local_min, value);
            local_max = fmaxf(local_max, value);
        };
        vectorized_process(threadIdx.x, blockDim.x, in_buf, previous_len, f);

        using BlockReduceT =
            hipcub::BlockReduce<T, BlockSize, hipcub::BLOCK_REDUCE_WARP_REDUCTIONS>;
        __shared__ typename BlockReduceT::TempStorage temp_storage;
        __shared__ bool use_one_pass;

        T global_min = BlockReduceT(temp_storage).Reduce(local_min, hipcub::Min());
        T global_max = BlockReduceT(temp_storage).Reduce(local_max, hipcub::Max());

        if(threadIdx.x == 0)
        {
            auto global_min_bits = twiddle_in(global_min, select_min);
            auto global_max_bits = twiddle_in(global_max, select_min);
            uint32_t diff        = global_min_bits ^ global_max_bits;
            use_one_pass         = diff < (1u << BitsPerPass);
        }
        __syncthreads();

        return use_one_pass;
    }
    else if(!out_buf)
    {
        // not use vectorized_process here because it increases #registers a lot
        auto const kth_value_bits    = counter->kth_value_bits;
        int const previous_start_bit = calc_start_bit<T, BitsPerPass>(pass - 1);

        for(IdxT i = threadIdx.x; i < previous_len; i += blockDim.x)
        {
            const T value            = in_buf[i];
            auto const previous_bits = (twiddle_in(value, select_min) >> previous_start_bit)
                                       << previous_start_bit;
            if(previous_bits == kth_value_bits)
            {
                int bucket = calc_bucket<T, BitsPerPass>(value, start_bit, mask, select_min);
                atomicAdd(histogram + bucket, static_cast<IdxT>(1));
            }
        }
    }
    else
    {
        // not use vectorized_process here because it increases #registers a lot
        IdxT* p_out_cnt              = &counter->out_cnt;
        auto const kth_value_bits    = counter->kth_value_bits;
        int const previous_start_bit = calc_start_bit<T, BitsPerPass>(pass - 1);

        if(in_idx_buf)
        {
            for(IdxT i = threadIdx.x; i < previous_len; i += blockDim.x)
            {
                const T value            = in_buf[i];
                auto const previous_bits = (twiddle_in(value, select_min) >> previous_start_bit)
                                           << previous_start_bit;
                if(previous_bits == kth_value_bits)
                {

                    IdxT pos         = atomicAdd(p_filter_cnt, static_cast<IdxT>(1));
                    out_buf[pos]     = value;
                    out_idx_buf[pos] = in_idx_buf[i];

                    int bucket = calc_bucket<T, BitsPerPass>(value, start_bit, mask, select_min);
                    atomicAdd(histogram + bucket, static_cast<IdxT>(1));
                }
                else if(previous_bits < kth_value_bits)
                {
                    IdxT pos = atomicAdd(p_out_cnt, static_cast<IdxT>(1));
                    if(WRITE_TOPK_VALUES)
                    {
                        out[pos] = value;
                    }
                    out_idx[pos] = in_idx_buf[i];
                }
            }
        }
        else
        {
            for(IdxT i = threadIdx.x; i < previous_len; i += blockDim.x)
            {
                const T value            = in_buf[i];
                auto const previous_bits = (twiddle_in(value, select_min) >> previous_start_bit)
                                           << previous_start_bit;
                if(previous_bits == kth_value_bits)
                {

                    IdxT pos         = atomicAdd(p_filter_cnt, static_cast<IdxT>(1));
                    out_buf[pos]     = value;
                    out_idx_buf[pos] = i;

                    int bucket = calc_bucket<T, BitsPerPass>(value, start_bit, mask, select_min);
                    atomicAdd(histogram + bucket, static_cast<IdxT>(1));
                }
                else if(previous_bits < kth_value_bits)
                {
                    IdxT pos = atomicAdd(p_out_cnt, static_cast<IdxT>(1));
                    if(WRITE_TOPK_VALUES)
                    {
                        out[pos] = value;
                    }
                    out_idx[pos] = i;
                }
            }
        }
    }

    return false;
}

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.