remove hybrid sort

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,47 +5,38 @@
 
 #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
-  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);
+  // 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);
   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);
@@ -59,15 +50,12 @@ 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 > kHybridSortMaxK) {
+  if (k > 64) {
     LaunchSort(topk_data, stream, scores_in, topk_data->intermediate_scores_1.get(),
                topk_data->intermediate_indices_1.get(), vocab_size, batch_size);
-  } else if (k <= 8 || vocab_size < 1024) {
+  } else {
     // 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,20 +18,18 @@ struct TopkData {
 
   // --- Intermediate Buffers for Top-K Algorithms ---
 
-  // Used to hold initial vocabulary indices for full sort, and intermediate
-  // indices during the reduction phase of hybrid sort.
+  // Used to hold initial vocabulary indices for full sort.
   cuda_unique_ptr<int> intermediate_indices_1;
 
-  // A dedicated "ping-pong" buffer for the hybrid sort index reduction.
+  // A dedicated "ping-pong" buffer for full sort.
   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 used by the hybrid sort's score reduction phase.
+  // A secondary "ping-pong" buffer.
   cuda_unique_ptr<float> intermediate_scores_2;
 
   // General-purpose temporary storage for CUB's DeviceSegmentedRadixSort (for full sort only).
@@ -44,7 +42,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, max_k (kHybridSortMaxK) for hybrid sort
+  int topk_stride = 0;                 // stride of the top-k output data: k for selection sort, vocab_size for full sort
 };
 
 // For parity test, a derived struct to help compact output buffers.
@@ -64,7 +62,6 @@ 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

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