submission 545370 · KernelIndex

Use it

Vendorable · source mirrored · license unknownView source →

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

sort_v6.py

curl "https://kernelindex.com/api/v1/implementations/kernelbot-sort-v2-545370?include=source"

interfacepython

Compatibility

measured onNVIDIA H100

declared hardwareNVIDIA H100

architecturessm_90

dtypesfp32

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload

Hardware

Latency

Rank

Observed

Sortsuite of 5 cases

NVIDIA H100

6.28ms

#9 of 26

2026-03-13

Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable

revision digestsha256:8aa75b6422510b66a65a1f0647f0a3c73f2a5239b8e6e006371474ee484bc805

license declaredunknown

license concludedunknown

authorsrajesh0042

imported2026-08-15

Kernel source

sort_v6.py31 lines

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

import torch
from task import input_t, output_t

# Sort: try stable=False + pre-allocate idx buffer + warmup
# The reference uses torch.sort(data)[0] which is stable=True by default
# But we just need matching sorted values, so stable=False is fine and might be faster

_idx_cache = {}

def _warmup():
    for size in [1024, 4096, 16384, 65536, 262144, 1048576]:
        x = torch.randn(size, device='cuda', dtype=torch.float32)
        idx = torch.empty(size, device='cuda', dtype=torch.int64)
        out = torch.empty_like(x)
        torch.sort(x, stable=False, out=(out, idx))
        torch.sort(x, stable=False, out=(out, idx))
    torch.cuda.synchronize()

_warmup()

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    n = data.numel()
    if n not in _idx_cache or _idx_cache[n].device != data.device:
        _idx_cache[n] = torch.empty(n, device=data.device, dtype=torch.int64)
    torch.sort(data, stable=False, out=(output, _idx_cache[n]))
    return output

scrolls · 31 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 545324.

  ⋯ 3 unchanged lines
  import torch
  from task import input_t, output_t
   
- # Sort with stable=False for potential speedup
- # The reference uses torch.sort(data)[0] which defaults to stable=False
- # Pre-allocate both values and indices buffers
+ # Sort: try stable=False + pre-allocate idx buffer + warmup
+ # The reference uses torch.sort(data)[0] which is stable=True by default
+ # But we just need matching sorted values, so stable=False is fine and might be faster
   
- _vals_cache = {}
  _idx_cache = {}
   
  def _warmup():
-     for size in [1024, 4096, 16384, 65536, 262144]:
+     for size in [1024, 4096, 16384, 65536, 262144, 1048576]:
          x = torch.randn(size, device='cuda', dtype=torch.float32)
-         torch.sort(x, stable=False)
+         idx = torch.empty(size, device='cuda', dtype=torch.int64)
+         out = torch.empty_like(x)
+         torch.sort(x, stable=False, out=(out, idx))
+         torch.sort(x, stable=False, out=(out, idx))
      torch.cuda.synchronize()
   
  _warmup()

scrolls · 26 diff lines total

Best evidence level for this revision: reported

JSON