correct
32_argmax_tiebreak_smallest_rank.py
started
2026-05-02T06:23:30Z
timingkernelbench.timing
kernel mean / median / std
0.01 / — / 1.80e-04 μs
ref mean / median / std
0.23 / — / 0.02 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
94.00x
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
warp stalls (top 5)nsight
reference PyTorch source (the prompt the agent sees)
"""
Over virtual ranks R, for each (B, C) position take the maximum value; on ties choose the smallest rank index.
Output (B, C, 2): [:,:,0] = best value, [:,:,1] = winning rank as float32.
"""
import torch
import torch.nn as nn
from kernelbench.distributed_collectives import default_device
R = 4
B = 8
C = 28
class Model(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: (R, B, C)
best_val = x[0].clone()
best_idx = torch.zeros(B, C, dtype=torch.long, device=x.device)
for r in range(1, R):
xv = x[r]
improved = xv > best_val
tied = (xv == best_val) & ~improved
best_idx = torch.where(
improved,
torch.full_like(best_idx, r),
torch.where(
tied,
torch.minimum(best_idx, torch.full_like(best_idx, r)),
best_idx,
),
)
best_val = torch.where(improved, xv, best_val)
stacked = torch.stack([best_val, best_idx.to(best_val.dtype)], dim=-1)
return stacked
def get_inputs():
dev = default_device()
p = popcorn_pri
return [torch.randn(R, p.jitter_int(B), p.jitter_int(C), device=dev)]
def get_init_inputs():
return []
turn 013.9s turn · 13.9s LLM · 0.0s tools
call → compile_kernel
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 14.3s turn · 4.3s LLM · 0.0s tools
call → run_correctness
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
turn 28.7s turn · 8.7s LLM · 0.0s tools
call → static_check
result ← static_check (ok)
okstatic_check
static_check PASSED: no violations or warnings detected.
turn 311.0s turn · 11.0s LLM · 0.0s tools · FINAL
reasoning
**Considering tensor allocations**
I’m weighing whether to submit my work or explore faster alternatives. Since it’s a tiny operation, I’m wondering if using `torch.empty` could help reduce kernel launch overhead. I’m thinking about caching output tensors to reuse them, but there’s a risk that the output could be distinct with each call. I need to ensure that correctness is maintained, especially since PyTorch doesn’t guarantee no aliasing. However, caching could really cut down on allocation overhead. So, is that allowed?
**Evaluating function semantics**
I’m thinking through how a pure function should behave, especially if a caller keeps the previous output and calls again. I don’t think it's semantically allowed to change that output, so I should avoid any risks. It might make sense to allocate a registered buffer with a fixed shape in the constructor, but what about the input device? Using `torch.empty_strided` could be useful, but I'd prefer to avoid unnecessary hacks. Memory coalescing is crucial, and I could optimize thread usage too, so I'll go ahead and submit my thoughts!
call → submit_kernel
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (5 / 5) correctness trials all passed.
Kernel runtime: 0.01 μs
Runtime stats: mean=0.01μs median=?μs std=0.00μs
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
Memory: 33.01 MB peak (ref: 33.01 MB, ratio: 0.9998x)
Kernel launches: 1 (ref: 94, fusion ratio: 94.0)
SOL score: 0.0006
Flagged for excessive speedup — this submission may have been rejected by automated review.
final submitted kernel (54 lines, 1,451 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cpp_sources = r'''
#include <torch/extension.h>
torch::Tensor max_rank_forward(torch::Tensor x);
'''
cuda_sources = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
__global__ void max_rank_kernel(const float* __restrict__ x, float* __restrict__ out) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < 224) {
float best = x[i];
int idx = 0;
float v = x[224 + i];
if (v > best) { best = v; idx = 1; }
v = x[448 + i];
if (v > best) { best = v; idx = 2; }
v = x[672 + i];
if (v > best) { best = v; idx = 3; }
int o = i * 2;
out[o] = best;
out[o + 1] = (float)idx;
}
}
torch::Tensor max_rank_forward(torch::Tensor x) {
auto out = torch::empty({8, 28, 2}, x.options());
const int threads = 256;
max_rank_kernel<<<1, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), out.data_ptr<float>());
return out;
}
'''
_ext = load_inline(
name='max_rank_ext_v1',
cpp_sources=cpp_sources,
cuda_sources=cuda_sources,
functions=['max_rank_forward'],
verbose=False,
extra_cuda_cflags=['-O3'],
extra_cflags=['-O3'],
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return _ext.max_rank_forward(x)