submission 682591

  ⋯ 6 unchanged lines
   
  cuda_src = r"""
  #include <torch/extension.h>
- #include <cuda_runtime.h>
  #include <cub/cub.cuh>
   
- // Pre-allocated temp storage
  static void* d_temp = nullptr;
  static size_t d_temp_size = 0;
   
- void sort_init(int64_t max_n) {
-     // Query temp storage size for max problem size
+ void sort_init(long long max_n) {
      size_t needed = 0;
+     // Use long long offset type for potentially faster CUB path
      cub::DeviceRadixSort::SortKeys(
          nullptr, needed,
          (const float*)nullptr, (float*)nullptr,
-         (int)max_n);
+         max_n, 0, 32);
   
      if (needed > d_temp_size) {
          if (d_temp) cudaFree(d_temp);
-         // Allocate with some headroom
-         d_temp_size = needed * 2;
+         d_temp_size = needed;
          cudaMalloc(&d_temp, d_temp_size);
      }
  }
   
- void sort_float(int64_t in_ptr, int64_t out_ptr, int N) {
+ void sort_raw(int64_t in_ptr, int64_t out_ptr, long long N) {
      size_t temp_size = d_temp_size;
      cub::DeviceRadixSort::SortKeys(
          d_temp, temp_size,
          reinterpret_cast<const float*>(in_ptr),
          reinterpret_cast<float*>(out_ptr),
-         N);
+         N, 0, 32);
  }
  """
   
  cpp_src = r"""
- void sort_init(int64_t max_n);
- void sort_float(int64_t in_ptr, int64_t out_ptr, int N);
+ void sort_init(long long max_n);
+ void sort_raw(int64_t in_ptr, int64_t out_ptr, long long N);
  """
   
  _ext = load_inline(
-     name="cub_radix_sort",
+     name="cub_sort_ll",
      cpp_sources=cpp_src,
      cuda_sources=cuda_src,
-     functions=["sort_init", "sort_float"],
+     functions=["sort_init", "sort_raw"],
      with_cuda=True,
      extra_cflags=["-O3"],
      extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_100a"],
      verbose=False,
  )
   
- # Pre-allocate temp storage for max size (100M elements)
+ # Pre-allocate for max size
  _ext.sort_init(100_000_000)
   
- # Warmup
- _wd = torch.randn(1000, device="cuda", dtype=torch.float32)
+ # Warmup with full size
+ _wd = torch.randn(100_000_000, device="cuda", dtype=torch.float32)
  _wo = torch.empty_like(_wd)
- _ext.sort_float(_wd.data_ptr(), _wo.data_ptr(), _wd.numel())
+ _ext.sort_raw(_wd.data_ptr(), _wo.data_ptr(), _wd.numel())
  torch.cuda.synchronize()
  del _wd, _wo
+ torch.cuda.empty_cache()
   
   
  def custom_kernel(data: input_t) -> output_t:
      data_in, output = data
-     _ext.sort_float(data_in.data_ptr(), output.data_ptr(), data_in.numel())
+     _ext.sort_raw(data_in.data_ptr(), output.data_ptr(), data_in.numel())
      return output