Revert "remove hybrid sort"

This reverts commit 55c4aa7.

Author: tianleiwu

src/cuda/cuda_sampling.h

@@ -42,4 +42,4 @@ void DispatchBlockwiseSoftmaxForward(cudaStream_t stream, float* output, const f
                                      int input_stride, int output_stride, int batch_count);
 
 }  // namespace cuda
-}  // namespace Generators
+}  // namespace Generators

src/cuda/cuda_topk.cu

@@ -5,38 +5,47 @@
 
 #include "cuda_topk.h"
 #include "cuda_topk_full_sort.cuh"
+#include "cuda_topk_hybrid_sort.cuh"
 #include "cuda_topk_select_sort.cuh"
 
 namespace Generators {
 namespace cuda {
 
+// Helper to determine the optimal partition size for the hybrid sort algorithm
+// based on vocabulary and batch size.
+inline int GetHybridSortPartitionSize(int vocab_size, int batch_size) {
+  if (vocab_size >= 147456) {
+    return (vocab_size > 256 * 1024) ? 8192 : 4096;
+  } else {
+    if (vocab_size >= 65536 || batch_size >= 4 && vocab_size > 49152) {
+      return 2048;
+    }
+  }
+
+  return 1024;
+}
+
 TopkData::TopkData(int batch_size, int vocab_size, cudaStream_t stream) {
+  // The intermediate buffers are used by hybrid and full sort algorithms.
+  int partition_size = GetHybridSortPartitionSize(vocab_size, batch_size);
+  size_t intermediate_buffer_elements = GetHybridSortIntermediateSize(batch_size, vocab_size, partition_size);
+
   size_t vocab_batch_size = static_cast<size_t>(vocab_size) * batch_size;
 
+  // Selection sort uses a buffer of batch_size * 64 elements, which is smaller than intermediate_buffer_elements.
+  size_t max_buffer_elements = std::max(vocab_batch_size, intermediate_buffer_elements);
+
   // Allocate all necessary device memory
-  // The buffers are sized for the full sort algorithm, which is the largest.
-  intermediate_indices_1 = CudaMallocArray<int>(vocab_batch_size);
-  intermediate_indices_2 = CudaMallocArray<int>(vocab_batch_size);
-  intermediate_scores_1 = CudaMallocArray<float>(vocab_batch_size);
-  intermediate_scores_2 = CudaMallocArray<float>(vocab_batch_size);
+  intermediate_indices_1 = CudaMallocArray<int>(max_buffer_elements);
+  intermediate_indices_2 = CudaMallocArray<int>(max_buffer_elements);
+  intermediate_scores_1 = CudaMallocArray<float>(max_buffer_elements);
+  intermediate_scores_2 = CudaMallocArray<float>(max_buffer_elements);
   batch_offsets = CudaMallocArray<int>(batch_size + 1);
 
   cub_temp_storage_bytes = GetFullSortCubTempStorageBytes(vocab_batch_size, batch_size, stream);
   cub_temp_storage = CudaMallocArray<unsigned char>(this->cub_temp_storage_bytes);
 }
 
-// Kernel to compact strided data into a dense layout.
-// Used to convert data from a [batch, stride] layout to a dense [batch, k] layout.
-template <typename T>
-__global__ void CompactStridedData(const T* input, T* output, int k, int batch_size, int input_stride) {
-  const int batch_idx = blockIdx.x;
-  for (int i = threadIdx.x; i < k; i += blockDim.x) {
-    int in_idx = batch_idx * input_stride + i;
-    int out_idx = batch_idx * k + i;
-    output[out_idx] = input[in_idx];
-  }
-}
-
 void TopkDataCompact::CompactOutput(int batch_size, int vocab_size, cudaStream_t stream, int k) {
   topk_scores_compact = CudaMallocArray<float>(static_cast<size_t>(batch_size) * k);
   topk_indices_compact = CudaMallocArray<int>(static_cast<size_t>(batch_size) * k);
@@ -50,12 +59,15 @@ void TopkDataCompact::CompactOutput(int batch_size, int vocab_size, cudaStream_t
 void GetTopK(TopkData* topk_data, cudaStream_t stream, const float* scores_in, int vocab_size, int batch_size, int k) {
   assert(topk_data != nullptr);
 
-  if (k > 64) {
+  if (k > kHybridSortMaxK) {
     LaunchSort(topk_data, stream, scores_in, topk_data->intermediate_scores_1.get(),
                topk_data->intermediate_indices_1.get(), vocab_size, batch_size);
-  } else {
+  } else if (k <= 8 || vocab_size < 1024) {
     // NOTE: This modifies scores_in in-place
     RunTopKViaSelectionSort(topk_data, stream, const_cast<float*>(scores_in), vocab_size, batch_size, k);
+  } else {
+    int partition_size = GetHybridSortPartitionSize(vocab_size, batch_size);
+    RunTopKViaHybridSort(topk_data, stream, scores_in, vocab_size, batch_size, k, partition_size);
   }
 }
 

src/cuda/cuda_topk.h

@@ -18,18 +18,20 @@ struct TopkData {
 
   // --- Intermediate Buffers for Top-K Algorithms ---
 
-  // Used to hold initial vocabulary indices for full sort.
+  // Used to hold initial vocabulary indices for full sort, and intermediate
+  // indices during the reduction phase of hybrid sort.
   cuda_unique_ptr<int> intermediate_indices_1;
 
-  // A dedicated "ping-pong" buffer for full sort.
+  // A dedicated "ping-pong" buffer for the hybrid sort index reduction.
   cuda_unique_ptr<int> intermediate_indices_2;
 
   // Primary buffer for holding raw scores.
   // - Full sort: Holds the fully sorted raw scores.
   // - Selection sort: Not used directly for output, but reserved.
+  // - Hybrid sort: Holds intermediate and final reduced raw scores.
   cuda_unique_ptr<float> intermediate_scores_1;
 
-  // A secondary "ping-pong" buffer.
+  // A secondary "ping-pong" buffer used by the hybrid sort's score reduction phase.
   cuda_unique_ptr<float> intermediate_scores_2;
 
   // General-purpose temporary storage for CUB's DeviceSegmentedRadixSort (for full sort only).
@@ -42,7 +44,7 @@ struct TopkData {
   // --- Information of Final Output (Input to Sampling Stage) ---
   const float* topk_scores = nullptr;  // pointer to the top-k scores data (in intermediate_scores_1 or intermediate_scores_2)
   const int* topk_indices = nullptr;   // pointer to the top-k indices data (in intermediate_indices_1 or intermediate_indices_2)
-  int topk_stride = 0;                 // stride of the top-k output data: k for selection sort, vocab_size for full sort
+  int topk_stride = 0;                 // stride of the top-k output data: k for selection sort, vocab_size for full sort, max_k (kHybridSortMaxK) for hybrid sort
 };
 
 // For parity test, a derived struct to help compact output buffers.
@@ -62,6 +64,7 @@ void GetTopK(TopkData* topk_data, cudaStream_t stream, const float* scores_in, i
 // to `scores_out` and `indices_out` in a compact [batch_size, k] layout.
 void RunTopKViaSelectionSort(TopkData* data, cudaStream_t stream, float* scores_in, int vocab_size, int batch_size, int k);
 void RunTopKViaFullSort(TopkData* data, cudaStream_t stream, const float* scores_in, int vocab_size, int batch_size, int k);
+void RunTopKViaHybridSort(TopkData* data, cudaStream_t stream, const float* scores_in, int vocab_size, int batch_size, int k, int partition_size);
 
 }  // namespace cuda
 }  // namespace Generators

src/cuda/cuda_topk_bitonic_sort_helper.cuh

@@ -0,0 +1,136 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <float.h>  // For FLT_MAX
+
+namespace Generators {
+namespace cuda {
+namespace bitonic {
+
+struct KeyValue {
+  float score;
+  int index;
+};
+
+// Performs a full bitonic sort in shared memory for `SortSize` elements.
+template <int kBlockSize, int SortSize>
+__device__ void SharedMemBitonicSort(KeyValue* smem_data) {
+  for (int k = 2; k <= SortSize; k <<= 1) {
+    for (int j = k >> 1; j > 0; j >>= 1) {
+      for (int i = threadIdx.x; i < SortSize; i += kBlockSize) {
+        int ixj = i ^ j;
+        if (ixj > i) {
+          bool ascending = ((i & k) == 0);
+          KeyValue* a = &smem_data[i];
+          KeyValue* b = &smem_data[ixj];
+          bool is_greater = (a->score > b->score) || (a->score == b->score && a->index < b->index);
+          if (is_greater != ascending) {
+            KeyValue temp = *a;
+            *a = *b;
+            *b = temp;
+          }
+        }
+      }
+      __syncthreads();
+    }
+  }
+
+  for (int j = SortSize >> 1; j > 0; j >>= 1) {
+    for (int i = threadIdx.x; i < SortSize; i += kBlockSize) {
+      int ixj = i ^ j;
+      if (ixj > i) {
+        KeyValue* a = &smem_data[i];
+        KeyValue* b = &smem_data[ixj];
+        if ((a->score < b->score) || (a->score == b->score && a->index > b->index)) {
+          KeyValue temp = *a;
+          *a = *b;
+          *b = temp;
+        }
+      }
+    }
+    __syncthreads();
+  }
+}
+
+namespace reduction {
+template <int N>
+__device__ void RegisterBitonicSort(float scores[N], int indices[N]) {
+  for (int k = 2; k <= N; k <<= 1) {
+    for (int j = k >> 1; j > 0; j >>= 1) {
+#pragma unroll
+      for (int i = 0; i < N; ++i) {
+        int ixj = i ^ j;
+        if (ixj > i) {
+          bool ascending = ((i & k) == 0);
+          bool is_greater = (scores[i] > scores[ixj]) || (scores[i] == scores[ixj] && indices[i] < indices[ixj]);
+          if (is_greater != ascending) {
+            float temp_s = scores[i];
+            scores[i] = scores[ixj];
+            scores[ixj] = temp_s;
+            int temp_i = indices[i];
+            indices[i] = indices[ixj];
+            indices[ixj] = temp_i;
+          }
+        }
+      }
+    }
+  }
+  for (int j = N >> 1; j > 0; j >>= 1) {
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      int ixj = i ^ j;
+      if (ixj > i) {
+        if ((scores[i] < scores[ixj]) || (scores[i] == scores[ixj] && indices[i] > indices[ixj])) {
+          float temp_s = scores[i];
+          scores[i] = scores[ixj];
+          scores[ixj] = temp_s;
+          int temp_i = indices[i];
+          indices[i] = indices[ixj];
+          indices[ixj] = temp_i;
+        }
+      }
+    }
+  }
+}
+
+template <int kBlockSize, int K, int PartitionsPerBlock>
+__global__ void BlockReduceTopK(const float* __restrict__ scores_in, const int* __restrict__ indices_in,
+                                float* __restrict__ scores_out, int* __restrict__ indices_out, int num_partitions_in) {
+  constexpr int SortSize = K * PartitionsPerBlock;
+  __shared__ KeyValue smem_buffer[SortSize];
+
+  const int batch_idx = blockIdx.y;
+  const int block_start_partition = blockIdx.x * PartitionsPerBlock;
+  const int num_partitions_to_process = min(PartitionsPerBlock, num_partitions_in - block_start_partition);
+
+  const int in_base_offset = batch_idx * num_partitions_in * K;
+  const int out_base_offset = (batch_idx * gridDim.x + blockIdx.x) * K;
+
+  for (int i = threadIdx.x; i < SortSize; i += kBlockSize) {
+    if (i < K * num_partitions_to_process) {
+      int partition_idx = i / K;
+      int element_idx = i % K;
+      int global_offset = in_base_offset + (block_start_partition + partition_idx) * K + element_idx;
+      smem_buffer[i].score = scores_in[global_offset];
+      smem_buffer[i].index = indices_in[global_offset];
+    } else {
+      smem_buffer[i].score = -FLT_MAX;
+      smem_buffer[i].index = -1;
+    }
+  }
+  __syncthreads();
+
+  SharedMemBitonicSort<kBlockSize, SortSize>(smem_buffer);
+
+  if (threadIdx.x < K) {
+    indices_out[out_base_offset + threadIdx.x] = smem_buffer[threadIdx.x].index;
+    scores_out[out_base_offset + threadIdx.x] = smem_buffer[threadIdx.x].score;
+  }
+}
+}  // namespace reduction
+
+}  // namespace bitonic
+}  // namespace cuda
+}  // namespace Generators

src/cuda/cuda_topk_full_sort.cuh

@@ -5,6 +5,7 @@
 
 #include <cub/device/device_segmented_radix_sort.cuh>
 #include "cuda_topk.h"
+#include "cuda_topk_hybrid_sort.cuh"  // For CompactStridedData
 
 namespace Generators {
 namespace cuda {