submission 67339

  ⋯ 2 unchanged lines
  from task import input_t, output_t
  import os
   
- # Ensure build directory exists
  os.makedirs("./cuda_build_sort", exist_ok=True)
   
- # Inline CUDA kernel using CUB's highly optimized radix sort
  cuda_source = """
  #include <torch/extension.h>
  #include <cuda_runtime.h>
  ⋯ 2 unchanged lines
   
  torch::Tensor cub_sort_kernel(torch::Tensor data, torch::Tensor output) {
      const int n = data.numel();
-  
-     // Get raw pointers
-     float* d_keys_in = data.data_ptr<float>();
-     float* d_keys_out = output.data_ptr<float>();
-  
-     // Get the CUDA stream
+     float* __restrict__ d_keys_in = data.data_ptr<float>();
+     float* __restrict__ d_keys_out = output.data_ptr<float>();
      cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
   
-     // Allocate temporary storage
+     int device;
+     cudaGetDevice(&device);
+     cudaDeviceProp prop;
+     cudaGetDeviceProperties(&prop, device);
+  
      void* d_temp_storage = nullptr;
      size_t temp_storage_bytes = 0;
   
-     // Determine temporary device storage requirements
+     int begin_bit = 0;
+     int end_bit = 32;
+  
+     // Get temp storage size
      cub::DeviceRadixSort::SortKeys(
          d_temp_storage,
          temp_storage_bytes,
          d_keys_in,
          d_keys_out,
          n,
-         0,              // begin_bit
-         32,             // end_bit (all 32 bits for float)
+         begin_bit,
+         end_bit,
          stream
      );
   
-     // Allocate temporary storage
-     cudaMalloc(&d_temp_storage, temp_storage_bytes);
+     // Architecture-specific allocation
+     if (prop.major == 9) {
+         // H100/B200: async allocation
+         cudaMallocAsync(&d_temp_storage, temp_storage_bytes, stream);
+     } else if (prop.major == 8 && prop.minor == 9) {
+         // L4: async allocation only (removed cudaMemAdvise - it was corrupting data)
+         cudaMallocAsync(&d_temp_storage, temp_storage_bytes, stream);
+     } else if (prop.major == 8 && prop.minor == 0) {
+         // A100: regular allocation
+         cudaMalloc(&d_temp_storage, temp_storage_bytes);
+     } else {
+         cudaMalloc(&d_temp_storage, temp_storage_bytes);
+     }
   
-     // Run sorting operation
+     // Execute sort
      cub::DeviceRadixSort::SortKeys(
          d_temp_storage,
          temp_storage_bytes,
          d_keys_in,
          d_keys_out,
          n,
-         0,              // begin_bit
-         32,             // end_bit
+         begin_bit,
+         end_bit,
          stream
      );
   
-     // Free temporary storage
-     cudaFree(d_temp_storage);
+     // Architecture-specific deallocation
+     if (prop.major == 9 || (prop.major == 8 && prop.minor == 9)) {
+         cudaFreeAsync(d_temp_storage, stream);
+     } else {
+         cudaFree(d_temp_storage);
+     }
   
      return output;
  }
  ⋯ 3 unchanged lines
  torch::Tensor cub_sort_kernel(torch::Tensor data, torch::Tensor output);
  """
   
- # Load the CUDA kernel with proper include paths
  try:
-     # Find CUDA toolkit path
      cuda_home = (
          os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") or "/usr/local/cuda"
      )
   
      cub_sort_module = load_inline(
-         name="cub_radix_sort",
+         name="cub_radix_sort_fixed",
          cpp_sources=cpp_source,
          cuda_sources=cuda_source,
          functions=["cub_sort_kernel"],
  ⋯ 3 unchanged lines
              "--use_fast_math",
              "-std=c++17",
              f"-I{cuda_home}/include",
+             "--expt-relaxed-constexpr",
+             "-lineinfo",
          ],
          extra_include_paths=[f"{cuda_home}/include"],
          build_directory="./cuda_build_sort",
-         verbose=True,
+         verbose=False,
      )
   
      def _custom_kernel(data: input_t) -> output_t:
-         """
-         Ultra-fast sort using CUB's DeviceRadixSort.
-         Args:
-             data: Tuple of (input_tensor, output_tensor)
-         Returns:
-             Sorted output tensor
-         """
          input_tensor, output_tensor = data
   
-         # Ensure contiguous memory layout for optimal performance
          if not input_tensor.is_contiguous():
              input_tensor = input_tensor.contiguous()
          if not output_tensor.is_contiguous():
              output_tensor = output_tensor.contiguous()
   
-         # Call CUB sort kernel
          cub_sort_module.cub_sort_kernel(input_tensor, output_tensor)
   
          return output_tensor
   
      custom_kernel = _custom_kernel
-     print("Successfully compiled CUB radix sort kernel!")
   
  except Exception as e:
-     print(f"Failed to compile CUB kernel: {e}")
-     print("Falling back to optimized torch.sort with torch.compile")
+     print(f"CUB kernel compilation failed: {e}")
   
-     # Fallback to optimized torch.sort
      def _custom_kernel(data: input_t) -> output_t:
          data, output = data
          output[...] = torch.sort(data)[0]