compile_fail
25_virtual_reduce_scatter_masked_sum.py
started
2026-05-11T10:13:50Z
compile error: Syntax error in custom generated code or ModelNew not found
timingkernelbench.timing
kernel mean / median / std
—
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 0306.5s turn · 13.7s LLM · 292.8s tools · tokens in=1,370 · out=1,434 · total=3,106 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model sums over the first dimension, masked by a boolean mask.
assistant```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void masked_sum_kernel(const float* __restrict__ x, const bool* __restrict__ mask, float* __restrict__ out, int R, int S) {
int j = blockIdx.x * blockDim.x + threadIdx.x;
if (j >= S) return;
float sum = 0.0f;
for (int r = 0; r < R; ++r) {
if (mask[r * S + j]) {
sum += x[r * S + j];
}
}
out[j] = sum;
}
torch::Tensor masked_sum_cuda(torch::Tensor x, torch::Tensor mask) {
TORCH_CHECK(x.is_cuda(), "x must be CUDA tensor");
TORCH_CHECK(mask.is_cuda(), "mask must be CUDA tensor");
auto S = x.size(1);
auto out = torch::zeros({S}, x.options());
const int threads = 256;
const int blocks = (S + threads - 1) / threads;
masked_sum_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
mask.data_ptr<bool>(),
out.data_ptr<float>(),
x.size(0),
S
);
return out;
}
'''
_cpp_source = r'''
torch::Tensor masked_sum_cuda(torch::Tensor x, torch::Tensor mask);
'''
_masked_sum_ext = load_inline(
name="masked_sum_ext",
cpp_sources=_cpp_source,
cuda_sources=_cuda_source,
functions=["masked_sum_cuda"],
verbose=False,
extra_cuda_cflags=["-O3"]
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
_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],
]
self.register_buffer(
"shard_mask",
torch.tensor(_MASK_DATA, dtype=torch.bool),
persistent=False,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
mask = self.shard_mask.to(device=x.device)
return _masked_sum_ext.masked_sum_cuda(x, mask)
```
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cuda_source = 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, const bool* __restrict__ mask, float* __restrict__ out, int R, int S) {\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n if (j >= S) return;\n float sum = 0.0f;\n for (int r = 0; r < R; ++r) {\n if (mask[r * S + j]) {\n sum += x[r * S + j];\n }\n }\n out[j] = sum;\n}\n\ntorch::Tensor masked_sum_cuda(torch::Tensor x, torch::Tensor mask) {\n TORCH_CHECK(x.is_cuda(), \"x must be CUDA tensor\");\n TORCH_CHECK(mask.is_cuda(), \"mask must be CUDA tensor\");\n auto S = x.size(1);\n auto out = torch::zeros({S}, x.options());\n const int threads = 256;\n const int blocks = (S + threads - 1) / threads;\n masked_sum_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n mask.data_ptr<bool>(),\n out.data_ptr<float>(),\n x.size(0),\n S\n );\n return out;\n}\n'''\n\n_cpp_source = r'''\ntorch::Tensor masked_sum_cuda(torch::Tensor x, torch::Tensor mask);\n'''\n\n_masked_sum_ext = load_inline(\n name=\"masked_sum_ext\",\n cpp_sources=_cpp_source,\n cuda_sources=_cuda_source,\n functions=[\"masked_sum_cuda\"],\n verbose=False,\n extra_cuda_cflags=[\"-O3\"]\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n _MASK_DATA = [\n [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],\n ]\n self.register_buffer(\n \"shard_mask\",\n torch.tensor(_MASK_DATA, dtype=torch.bool),\n persistent=False,\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n mask = self.shard_mask.to(device=x.device)\n return _masked_sum_ext.masked_sum_cuda(x, mask)\n```"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: kernel did not compile.
Syntax error in custom generated code or ModelNew not found
final submitted kernel (71 lines, 2,101 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void masked_sum_kernel(const float* __restrict__ x, const bool* __restrict__ mask, float* __restrict__ out, int R, int S) {
int j = blockIdx.x * blockDim.x + threadIdx.x;
if (j >= S) return;
float sum = 0.0f;
for (int r = 0; r < R; ++r) {
if (mask[r * S + j]) {
sum += x[r * S + j];
}
}
out[j] = sum;
}
torch::Tensor masked_sum_cuda(torch::Tensor x, torch::Tensor mask) {
TORCH_CHECK(x.is_cuda(), "x must be CUDA tensor");
TORCH_CHECK(mask.is_cuda(), "mask must be CUDA tensor");
auto S = x.size(1);
auto out = torch::zeros({S}, x.options());
const int threads = 256;
const int blocks = (S + threads - 1) / threads;
masked_sum_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
mask.data_ptr<bool>(),
out.data_ptr<float>(),
x.size(0),
S
);
return out;
}
'''
_cpp_source = r'''
torch::Tensor masked_sum_cuda(torch::Tensor x, torch::Tensor mask);
'''
_masked_sum_ext = load_inline(
name="masked_sum_ext",
cpp_sources=_cpp_source,
cuda_sources=_cuda_source,
functions=["masked_sum_cuda"],
verbose=False,
extra_cuda_cflags=["-O3"]
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
_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],
]
self.register_buffer(
"shard_mask",
torch.tensor(_MASK_DATA, dtype=torch.bool),
persistent=False,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
mask = self.shard_mask.to(device=x.device)
return _masked_sum_ext.masked_sum_cuda(x, mask)
```