compile_fail
25_virtual_reduce_scatter_masked_sum.py
started
2026-05-11T22:36:58Z
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 015.0s turn · 15.0s LLM · 0.0s tools · tokens in=2,268 · out=501 · total=3,315
reasoning
The task is to replace the forward computation with a custom CUDA kernel. The reference is:
The task requires replacing the forward computation of a PyTorch model with a custom CUDA kernel. The original model computes a masked sum on a small tensor of shape (4,20), multiplying an input tensor by a fixed boolean mask and summing along the first dimension.
assistant```python
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 masked_sum_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int S) {
int j = blockIdx.x * blockDim.x + threadIdx.x;
if (j >= S) return;
float sum = 0.0f;
// Hardcoded mask from _MASK_DATA
// Column masks: 0-7:4, 8-9:3, 10-11:2, 12-17:1, 18-19:0
if (j <= 7) {
sum = x[0*S + j] + x[1*S + j] + x[2*S + j] + x[3*S + j];
} else if (j <= 9) {
sum = x[1*S + j] + x[2*S + j] + x[3*S + j];
} else if (j <= 11) {
sum = x[1*S + j] + x[3*S + j];
} else if (j <= 17) {
sum = x[3*S + j];
}
out[j] = sum;
}
torch::Tensor masked_sum_cuda(torch::Tensor x) {
auto S = x.size(1);
auto out = torch::empty({S}, x.options());
const int threads = 32;
const int blocks = (S + threads - 1) / threads;
masked_sum_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
out.data_ptr<float>(),
x.size(0),
S
);
return out;
}
'''
cpp_source = '''
torch::Tensor masked_sum_cuda(torch::Tensor x);
'''
_masked_sum = load_inline(
name='masked_sum',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['masked_sum_cuda'],
verbose=False,
with_cuda=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return _masked_sum.masked_sum_cuda(x)
```
turn 181.7s turn · 6.0s LLM · 75.7s tools · tokens in=3,383 · out=508 · total=3,991
reasoning
The instruction is: "You must call at least one tool this turn: use `compile_kernel` / `run_correctness` to iterate, or call `submit_kernel` with your best `ModelNew` implementation to finalize the ru...
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 masked_sum_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int S) {\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n if (j >= S) return;\n \n float sum = 0.0f;\n // Hardcoded mask from _MASK_DATA\n // Column masks: 0-7:4, 8-9:3, 10-11:2, 12-17:1, 18-19:0\n if (j <= 7) {\n sum = x[0*S + j] + x[1*S + j] + x[2*S + j] + x[3*S + j];\n } else if (j <= 9) {\n sum = x[1*S + j] + x[2*S + j] + x[3*S + j];\n } else if (j <= 11) {\n sum = x[1*S + j] + x[3*S + j];\n } else if (j <= 17) {\n sum = x[3*S + j];\n }\n out[j] = sum;\n}\n\ntorch::Tensor masked_sum_cuda(torch::Tensor x) {\n auto S = x.size(1);\n auto out = torch::empty({S}, x.options());\n \n const int threads = 32;\n const int blocks = (S + threads - 1) / threads;\n \n masked_sum_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n out.data_ptr<float>(),\n x.size(0),\n S\n );\n \n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor masked_sum_cuda(torch::Tensor x);\n'''\n\n_masked_sum = load_inline(\n name='masked_sum',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['masked_sum_cuda'],\n verbose=False,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return _masked_sum.masked_sum_cuda(x)\n"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 23.6s turn · 3.5s LLM · 0.0s tools · tokens in=4,008 · out=509 · total=4,531
reasoning
It compiled successfully. Now I should check correctness with run_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 masked_sum_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int S) {\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n if (j >= S) return;\n \n float sum = 0.0f;\n // Hardcoded mask from _MASK_DATA\n // Column masks: 0-7:4, 8-9:3, 10-11:2, 12-17:1, 18-19:0\n if (j <= 7) {\n sum = x[0*S + j] + x[1*S + j] + x[2*S + j] + x[3*S + j];\n } else if (j <= 9) {\n sum = x[1*S + j] + x[2*S + j] + x[3*S + j];\n } else if (j <= 11) {\n sum = x[1*S + j] + x[3*S + j];\n } else if (j <= 17) {\n sum = x[3*S + j];\n }\n out[j] = sum;\n}\n\ntorch::Tensor masked_sum_cuda(torch::Tensor x) {\n auto S = x.size(1);\n auto out = torch::empty({S}, x.options());\n \n const int threads = 32;\n const int blocks = (S + threads - 1) / threads;\n \n masked_sum_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n out.data_ptr<float>(),\n x.size(0),\n S\n );\n \n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor masked_sum_cuda(torch::Tensor x);\n'''\n\n_masked_sum = load_inline(\n name='masked_sum',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['masked_sum_cuda'],\n verbose=False,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return _masked_sum.masked_sum_cuda(x)\n"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (3 / 3) 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 36.7s turn · 6.7s LLM · 0.0s tools · tokens in=4,591 · out=509 · total=5,148 · FINAL
reasoning
It passed correctness. Now, since it's tiny, it should be faster or at least I can submit it.
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 masked_sum_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int S) {\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n if (j >= S) return;\n \n float sum = 0.0f;\n // Hardcoded mask from _MASK_DATA\n // Column masks: 0-7:4, 8-9:3, 10-11:2, 12-17:1, 18-19:0\n if (j <= 7) {\n sum = x[0*S + j] + x[1*S + j] + x[2*S + j] + x[3*S + j];\n } else if (j <= 9) {\n sum = x[1*S + j] + x[2*S + j] + x[3*S + j];\n } else if (j <= 11) {\n sum = x[1*S + j] + x[3*S + j];\n } else if (j <= 17) {\n sum = x[3*S + j];\n }\n out[j] = sum;\n}\n\ntorch::Tensor masked_sum_cuda(torch::Tensor x) {\n auto S = x.size(1);\n auto out = torch::empty({S}, x.options());\n \n const int threads = 32;\n const int blocks = (S + threads - 1) / threads;\n \n masked_sum_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n out.data_ptr<float>(),\n x.size(0),\n S\n );\n \n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor masked_sum_cuda(torch::Tensor x);\n'''\n\n_masked_sum = load_inline(\n name='masked_sum',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['masked_sum_cuda'],\n verbose=False,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return _masked_sum.masked_sum_cuda(x)\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 (66 lines, 1,611 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 masked_sum_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int S) {
int j = blockIdx.x * blockDim.x + threadIdx.x;
if (j >= S) return;
float sum = 0.0f;
// Hardcoded mask from _MASK_DATA
// Column masks: 0-7:4, 8-9:3, 10-11:2, 12-17:1, 18-19:0
if (j <= 7) {
sum = x[0*S + j] + x[1*S + j] + x[2*S + j] + x[3*S + j];
} else if (j <= 9) {
sum = x[1*S + j] + x[2*S + j] + x[3*S + j];
} else if (j <= 11) {
sum = x[1*S + j] + x[3*S + j];
} else if (j <= 17) {
sum = x[3*S + j];
}
out[j] = sum;
}
torch::Tensor masked_sum_cuda(torch::Tensor x) {
auto S = x.size(1);
auto out = torch::empty({S}, x.options());
const int threads = 32;
const int blocks = (S + threads - 1) / threads;
masked_sum_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
out.data_ptr<float>(),
x.size(0),
S
);
return out;
}
'''
cpp_source = '''
torch::Tensor masked_sum_cuda(torch::Tensor x);
'''
_masked_sum = load_inline(
name='masked_sum',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['masked_sum_cuda'],
verbose=False,
with_cuda=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return _masked_sum.masked_sum_cuda(x)
```