correct
32_argmax_tiebreak_smallest_rank.py
started
2026-05-11T09:19:31Z
timingkernelbench.timing
kernel mean / median / std
0.01 / — / 1.56e-04 μs
ref mean / median / std
0.28 / — / 0.03 μ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 0294.5s turn · 39.2s LLM · 255.3s tools · tokens in=1,207 · out=476 · total=2,467 · FINAL
reasoning
The task is to replace the forward computation with a custom CUDA kernel. The reference is a PyTorch model that finds the maximum value over virtual ranks R for each (B, C) position, and on ties, choo...I need to implement a custom CUDA kernel to replace the forward computation in a PyTorch model. The model processes an input tensor x of shape (R, B, C), where R is the number of ranks, B is batch size, and C is channels. For each (B, C) position, it finds 1) the maximum value across the R ranks, and 2) if there are ties, the smallest rank index. The output is a tensor of shape (B, C, 2) containing the best value and the winning rank as floats.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n\n__global__ void max_over_ranks_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int C) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * C;\n if (idx < total) {\n float best_val = x[0 * total + idx];\n int best_r = 0;\n for (int r = 1; r < R; ++r) {\n float val = x[r * total + idx];\n if (val > best_val || (val == best_val && r < best_r)) {\n best_val = val;\n best_r = r;\n }\n }\n out[idx * 2 + 0] = best_val;\n out[idx * 2 + 1] = (float)best_r;\n }\n}\n\ntorch::Tensor max_over_ranks(torch::Tensor x) {\n TORCH_CHECK(x.dim() == 3, \"x must be 3D\");\n auto R = x.size(0);\n auto B = x.size(1);\n auto C = x.size(2);\n auto out = torch::empty({B, C, 2}, x.options());\n int total = B * C;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n max_over_ranks_kernel<<<blocks, threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), R, B, C);\n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor max_over_ranks(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='max_ranks_ext',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['max_over_ranks'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.op = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.op.max_over_ranks(x)\n"
}
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.001
Flagged for excessive speedup — this submission may have been rejected by automated review.
final submitted kernel (60 lines, 1,637 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void max_over_ranks_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int C) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * C;
if (idx < total) {
float best_val = x[0 * total + idx];
int best_r = 0;
for (int r = 1; r < R; ++r) {
float val = x[r * total + idx];
if (val > best_val || (val == best_val && r < best_r)) {
best_val = val;
best_r = r;
}
}
out[idx * 2 + 0] = best_val;
out[idx * 2 + 1] = (float)best_r;
}
}
torch::Tensor max_over_ranks(torch::Tensor x) {
TORCH_CHECK(x.dim() == 3, "x must be 3D");
auto R = x.size(0);
auto B = x.size(1);
auto C = x.size(2);
auto out = torch::empty({B, C, 2}, x.options());
int total = B * C;
int threads = 256;
int blocks = (total + threads - 1) / threads;
max_over_ranks_kernel<<<blocks, threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), R, B, C);
return out;
}
'''
cpp_source = '''
torch::Tensor max_over_ranks(torch::Tensor x);
'''
module = load_inline(
name='max_ranks_ext',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['max_over_ranks'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.op = module
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.op.max_over_ranks(x)