correct
32_argmax_tiebreak_smallest_rank.py
started
2026-05-11T11:54:09Z
timingkernelbench.timing
kernel mean / median / std
0.01 / — / 0.02 μs
ref mean / median / std
0.25 / — / 0.07 μ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 0181.7s turn · 14.5s LLM · 167.2s tools · tokens in=1,574 · out=507 · total=2,778
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model finds, over virtual ranks R, for each (B, C) position, the maximum value, and on ties, chooses th...I need to replace the forward computation of a PyTorch model with a custom CUDA kernel in a new ModelNew class. The input is a tensor x of shape (R, B, dim=-1), but practically (4, 8, 28), and I need to output a stacked tensor of shape (B dim=-1, 2), where for each (B, C) position, I find the maximum value across the R ranks, and on ties, the smallest rank index. The output has the best value and the winning rank as float32.
call → compile_kernel
callcompile_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.h>\n#include <cuda_runtime.h>\n\n__global__ void max_min_rank_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 int b = idx / C;\n int c = idx % C;\n float best_val = x[b * C + c];\n int best_r = 0;\n for (int r = 1; r < R; r++) {\n float val = x[r * B * C + b * C + c];\n if (val > best_val) {\n best_val = val;\n best_r = r;\n } else if (val == best_val && r < best_r) {\n best_r = r;\n }\n }\n out[(b * C + c) * 2 + 0] = best_val;\n out[(b * C + c) * 2 + 1] = (float)best_r;\n }\n}\n\ntorch::Tensor max_min_rank_cuda(torch::Tensor x) {\n int R = x.size(0);\n int B = x.size(1);\n int 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_min_rank_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_min_rank_cuda(torch::Tensor x);\n\"\"\"\n\nmodule = load_inline(\n name=\"max_min_rank\",\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=[\"max_min_rank_cuda\"],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.module.max_min_rank_cuda(x)\n"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 1164.0s turn · 4.3s LLM · 159.7s tools · tokens in=2,795 · out=508 · total=3,310
reasoning
Compiled successfully. Now check correctness.
call → run_correctness
callrun_correctness
{
"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.h>\n#include <cuda_runtime.h>\n\n__global__ void max_min_rank_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 int b = idx / C;\n int c = idx % C;\n float best_val = x[b * C + c];\n int best_r = 0;\n for (int r = 1; r < R; r++) {\n float val = x[r * B * C + b * C + c];\n if (val > best_val) {\n best_val = val;\n best_r = r;\n } else if (val == best_val && r < best_r) {\n best_r = r;\n }\n }\n out[(b * C + c) * 2 + 0] = best_val;\n out[(b * C + c) * 2 + 1] = (float)best_r;\n }\n}\n\ntorch::Tensor max_min_rank_cuda(torch::Tensor x) {\n int R = x.size(0);\n int B = x.size(1);\n int 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_min_rank_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_min_rank_cuda(torch::Tensor x);\n\"\"\"\n\nmodule = load_inline(\n name=\"max_min_rank\",\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=[\"max_min_rank_cuda\"],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.module.max_min_rank_cuda(x)\n"
}
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 2233.3s turn · 5.1s LLM · 228.2s tools · tokens in=3,370 · out=507 · total=3,933 · FINAL
reasoning
Correct. Now, to submit, since it's tiny, probably fast enough.
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.h>\n#include <cuda_runtime.h>\n\n__global__ void max_min_rank_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 int b = idx / C;\n int c = idx % C;\n float best_val = x[b * C + c];\n int best_r = 0;\n for (int r = 1; r < R; r++) {\n float val = x[r * B * C + b * C + c];\n if (val > best_val) {\n best_val = val;\n best_r = r;\n } else if (val == best_val && r < best_r) {\n best_r = r;\n }\n }\n out[(b * C + c) * 2 + 0] = best_val;\n out[(b * C + c) * 2 + 1] = (float)best_r;\n }\n}\n\ntorch::Tensor max_min_rank_cuda(torch::Tensor x) {\n int R = x.size(0);\n int B = x.size(1);\n int 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_min_rank_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_min_rank_cuda(torch::Tensor x);\n\"\"\"\n\nmodule = load_inline(\n name=\"max_min_rank\",\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=[\"max_min_rank_cuda\"],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.module.max_min_rank_cuda(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.02μ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.0009
Flagged for excessive speedup — this submission may have been rejected by automated review.
final submitted kernel (65 lines, 1,758 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.h>
#include <cuda_runtime.h>
__global__ void max_min_rank_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) {
int b = idx / C;
int c = idx % C;
float best_val = x[b * C + c];
int best_r = 0;
for (int r = 1; r < R; r++) {
float val = x[r * B * C + b * C + c];
if (val > best_val) {
best_val = val;
best_r = r;
} else if (val == best_val && r < best_r) {
best_r = r;
}
}
out[(b * C + c) * 2 + 0] = best_val;
out[(b * C + c) * 2 + 1] = (float)best_r;
}
}
torch::Tensor max_min_rank_cuda(torch::Tensor x) {
int R = x.size(0);
int B = x.size(1);
int 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_min_rank_kernel<<<blocks, threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), R, B, C);
return out;
}
"""
cpp_source = """
torch::Tensor max_min_rank_cuda(torch::Tensor x);
"""
module = load_inline(
name="max_min_rank",
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=["max_min_rank_cuda"],
verbose=True,
with_cuda=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.module = module
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.module.max_min_rank_cuda(x)