Skip to content
KernelIndex
Search⌘K
Sign in

submission 612537

dannywillowliu-uchi · python · License unknown
Kernel source · 53 lines ↓

Use it

Vendorable · source mirrored · license unknownView source →

No package. Vendor the mirrored source: 53 lines, June 9 Researcher Reciprocity License v1.0.

submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-sort-v2-612537?include=source"
interfacepython
Compatibility
measured onNVIDIA B200
declared hardwareNVIDIA B200
architecturessm_100
dtypesfp32

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload
Hardware
Latency
Rank
Observed
Sortsuite of 5 cases
NVIDIA B200
2.04ms
#7 of 23
2026-03-23
Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable
revision digestsha256:7ddb69fd7fb7dd22d1363002407fb3a35fcb442708c7d5789648b1234b1e24d5
license declaredunknown
license concludedunknown
authorsdannywillowliu-uchi
imported2026-08-15

Kernel source

submission.py53 lines
import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

cuda_source = r"""
#include <torch/extension.h>
#include <cub/cub.cuh>

static void* g_temp = nullptr;
static size_t g_temp_bytes = 0;

torch::Tensor cub_radix_sort(torch::Tensor input, torch::Tensor output) {
	const int n = input.numel();
	unsigned int* d_in = reinterpret_cast<unsigned int*>(input.data_ptr<float>());
	unsigned int* d_out = reinterpret_cast<unsigned int*>(output.data_ptr<float>());

	void* d_temp = nullptr;
	size_t temp_bytes = 0;
	cub::DeviceRadixSort::SortKeys(d_temp, temp_bytes, d_in, d_out, n);

	if (temp_bytes > g_temp_bytes) {
		if (g_temp) cudaFree(g_temp);
		cudaMalloc(&g_temp, temp_bytes);
		g_temp_bytes = temp_bytes;
	}

	cub::DeviceRadixSort::SortKeys(g_temp, g_temp_bytes, d_in, d_out, n);
	return output;
}
"""

cpp_source = """
torch::Tensor cub_radix_sort(torch::Tensor input, torch::Tensor output);
"""

module = load_inline(
	name="cub_radix_sort_fast",
	cpp_sources=[cpp_source],
	cuda_sources=[cuda_source],
	functions=["cub_radix_sort"],
	extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
	verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
	inp, out = data
	module.cub_radix_sort(inp, out)
	return out
scrolls · 53 lines total

Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0

Changes from previous submission

Against this author's previous submission submission 612483.

  ⋯ 4 unchanged lines
  from torch.utils.cpp_extension import load_inline
  from task import input_t, output_t
   
- cuda_source = """
+ cuda_source = r"""
  #include <torch/extension.h>
  #include <cub/cub.cuh>
   
- // Persistent temp storage to avoid re-allocation
  static void* g_temp = nullptr;
  static size_t g_temp_bytes = 0;
   
  torch::Tensor cub_radix_sort(torch::Tensor input, torch::Tensor output) {
  	const int n = input.numel();
- 	float* d_in = input.data_ptr<float>();
- 	float* d_out = output.data_ptr<float>();
+ 	unsigned int* d_in = reinterpret_cast<unsigned int*>(input.data_ptr<float>());
+ 	unsigned int* d_out = reinterpret_cast<unsigned int*>(output.data_ptr<float>());
   
- 	// Determine temp storage needed
  	void* d_temp = nullptr;
  	size_t temp_bytes = 0;
  	cub::DeviceRadixSort::SortKeys(d_temp, temp_bytes, d_in, d_out, n);
   
- 	// Re-allocate only if needed
  	if (temp_bytes > g_temp_bytes) {
  		if (g_temp) cudaFree(g_temp);
  		cudaMalloc(&g_temp, temp_bytes);
  ⋯ 1 unchanged lines
  	}
   
  	cub::DeviceRadixSort::SortKeys(g_temp, g_temp_bytes, d_in, d_out, n);
-  
  	return output;
  }
  """
  ⋯ 3 unchanged lines
  """
   
  module = load_inline(
- 	name="cub_radix_sort",
+ 	name="cub_radix_sort_fast",
  	cpp_sources=[cpp_source],
  	cuda_sources=[cuda_source],
  	functions=["cub_radix_sort"],
scrolls · 46 diff lines total

Best evidence level for this revision: reported

JSON