reference PyTorch source (the prompt the agent sees)
"""
Unequal virtual shard lengths: x is (R, S_MAX) with a fixed boolean mask marking valid cells.
Output (S_MAX,) is sum_r x[r, j] * mask[r, j] (masked sum along the virtual rank dimension).
Padding positions are masked out so they do not contribute.
"""
import torch
import torch.nn as nn
from kernelbench.distributed_collectives import default_device
R = 4
S_MAX = 20
# mask[r, j] == True iff rank r contributes a real value at column j
_MASK_DATA = [
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
]
class Model(nn.Module):
def __init__(self):
super().__init__()
self.register_buffer(
"shard_mask",
torch.tensor(_MASK_DATA, dtype=torch.bool),
persistent=False,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
m = self.shard_mask.to(dtype=x.dtype, device=x.device)
return (x * m).sum(dim=0)
def get_inputs():
dev = default_device()
p = popcorn_pri
return [torch.randn(p.jitter_int(R), p.jitter_int(S_MAX), device=dev)]
def get_init_inputs():
return []
turn 012.4s turn · 12.4s LLM · 0.0s tools
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r\"\"\"\n#include <torch/extension.h>\n\ntorch::Tensor masked_sum_forward(torch::Tensor x);\n\"\"\"\n\n_cuda_src = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void masked_sum_kernel(const float* __restrict__ x, float* __restrict__ y) {\n int j = threadIdx.x;\n if (j < 20) {\n float v = 0.0f;\n // mask rows: r0 valid 0..7, r1 valid 0..11, r2 valid 0..9, r3 valid 0..17\n if (j < 8) v += x[j];\n if (j < 12) v += x[20 + j];\n if (j < 10) v += x[40 + j];\n if (j < 18) v += x[60 + j];\n y[j] = v;\n }\n}\n\ntorch::Tensor masked_sum_forward(torch::Tensor x) {\n auto y = torch::empty({20}, x.options());\n masked_sum_kernel<<<1, 32, 0, at::cuda::getCurrentCUDAStream()>>>(\n reinterpret_cast<const float*>(x.data_ptr<float>()),\n reinterpret_cast<float*>(y.data_ptr<float>()));\n return y;\n}\n\"\"\"\n\n_ext = load_inline(\n name=\"masked_sum_ext_v1\",\n cpp_sources=_cpp_src,\n cuda_sources=_cuda_src,\n functions=[\"masked_sum_forward\"],\n extra_cflags=[\"-O3\"],\n extra_cuda_cflags=[\"-O3\"],\n verbose=False,\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return _ext.masked_sum_forward(x)\n"
}
result ← compile_kernel (fail)
failcompile_kernel
compile_kernel FAILED: unexpected error during tool execution.
_WorkItemTimeout: work item exceeded soft timeout budget
Traceback (most recent call last):
File "/scratch/tejas/PopcornBench/src/kernelbench/agent/agent.py", line 627, in _execute_tool
return tool.execute(self.ctx, **args)
File "/scratch/tejas/PopcornBench/src/kernelbench/agent/tools.py", line 292, in execute
return ctx.eval_client.compile_kernel(ctx, kernel_code)
File "/scratch/tejas/PopcornBench/src/kernelbench/agent/eval_client.py", line 54, in compile_kernel
return self._call("compile", args)
File "/scratch/tejas/PopcornBench/src/kernelbench/agent/eval_client.py", line 117, in _call
payload = self._response_q.get(timeout=remaining)
File "<string>", line 2, in get
File "/home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/multiprocessing/managers.py", line 818, in _callmethod
kind, result = conn.recv()
File "/home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/multiprocessing/connection.py", line 250, in recv
buf = self._recv_bytes()
File "/home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/multiprocessing/connection.py", line 414, in _recv_bytes
buf = self._recv(4)
File "/home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/multiprocessing/connection.py", line 379, in _recv
chunk = read(handle, remaining)
File "/scratch/tejas/PopcornBench/scripts/run_sweep.py", line 161, in _alarm_handler
raise _WorkItemTimeout("work item exceeded soft timeout budget")
__mp_main__._WorkItemTimeout: work item exceeded soft timeout budget