compile_fail
49_sum_gelu_across_virtual_ranks.py
started
2026-05-11T09:25:17Z
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)
"""
Order-sensitive virtual collective: out = sum_r GELU(x_r), same shape as one rank slice (B, D).
Note: this is not equal to GELU(sum_r x_r).
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from kernelbench.distributed_collectives import default_device
R = 4
B = 14
D = 40
class Model(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.gelu(x).sum(dim=0)
def get_inputs():
dev = default_device()
p = popcorn_pri
return [torch.randn(p.jitter_int(R), p.jitter_int(B), p.jitter_int(D, align=8), device=dev)]
def get_init_inputs():
return []
turn 025.5s turn · 25.5s LLM · 0.0s tools · tokens in=1,048 · out=1,201 · total=2,738 · FINAL
reasoning
The task is to replace the forward computation with a custom CUDA kernel for this model:
I need to replace the forward computation with a custom CUDA kernel for the model. The model is a simple PyTorch module that applies GELU to the input tensor and sums over the first dimension.
assistant```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__device__ __forceinline__ float gelu(float x) {
const float kAlpha = 0.7978845608028654f; // sqrt(2/pi)
const float kBeta = 0.044715f;
float x3 = x * x * x;
float inner = kAlpha * (x + kBeta * x3);
return 0.5f * x * (1.0f + tanhf(inner));
}
__global__ void fused_gelu_sum_kernel(const float* __restrict__ input, float* __restrict__ output,
int R, int B, int D) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * D;
if (idx >= total) return;
int b = idx / D;
int d = idx % D;
float sum = 0.0f;
for (int r = 0; r < R; r++) {
float val = input[r * B * D + b * D + d];
sum += gelu(val);
}
output[b * D + d] = sum;
}
torch::Tensor fused_gelu_sum_cuda(torch::Tensor x) {
TORCH_CHECK(x.is_cuda(), "x must be CUDA tensor");
TORCH_CHECK(x.dim() == 3, "x must be 3D");
int R = x.size(0);
int B = x.size(1);
int D = x.size(2);
auto output = torch::empty({B, D}, x.options());
int total = B * D;
int threads = 256;
int blocks = (total + threads - 1) / threads;
fused_gelu_sum_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
output.data_ptr<float>(),
R, B, D
);
return output;
}
'''
cpp_source = r'''
torch::Tensor fused_gelu_sum_cuda(torch::Tensor x);
'''
# Compile the extension
_mod = load_inline(
name="fused_gelu_sum",
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=["fused_gelu_sum_cuda"],
verbose=False,
extra_cuda_cflags=["-O3"]
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return _mod.fused_gelu_sum_cuda(x)
``` submit_kernel with kernel_code is ... (full code above) (wait, I must do the actual call)
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__device__ __forceinline__ float gelu(float x) {\n const float kAlpha = 0.7978845608028654f; // sqrt(2/pi)\n const float kBeta = 0.044715f;\n float x3 = x * x * x;\n float inner = kAlpha * (x + kBeta * x3);\n return 0.5f * x * (1.0f + tanhf(inner));\n}\n\n__global__ void fused_gelu_sum_kernel(const float* __restrict__ input, float* __restrict__ output,\n int R, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * D;\n if (idx >= total) return;\n \n int b = idx / D;\n int d = idx % D;\n \n float sum = 0.0f;\n for (int r = 0; r < R; r++) {\n float val = input[r * B * D + b * D + d];\n sum += gelu(val);\n }\n output[b * D + d] = sum;\n}\n\ntorch::Tensor fused_gelu_sum_cuda(torch::Tensor x) {\n TORCH_CHECK(x.is_cuda(), \"x must be CUDA tensor\");\n TORCH_CHECK(x.dim() == 3, \"x must be 3D\");\n \n int R = x.size(0);\n int B = x.size(1);\n int D = x.size(2);\n \n auto output = torch::empty({B, D}, x.options());\n \n int total = B * D;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n \n fused_gelu_sum_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n output.data_ptr<float>(),\n R, B, D\n );\n \n return output;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor fused_gelu_sum_cuda(torch::Tensor x);\n'''\n\n# Compile the extension\n_mod = load_inline(\n name=\"fused_gelu_sum\",\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=[\"fused_gelu_sum_cuda\"],\n verbose=False,\n extra_cuda_cflags=[\"-O3\"]\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 _mod.fused_gelu_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 (80 lines, 2,005 chars) — full source, untruncated
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__device__ __forceinline__ float gelu(float x) {
const float kAlpha = 0.7978845608028654f; // sqrt(2/pi)
const float kBeta = 0.044715f;
float x3 = x * x * x;
float inner = kAlpha * (x + kBeta * x3);
return 0.5f * x * (1.0f + tanhf(inner));
}
__global__ void fused_gelu_sum_kernel(const float* __restrict__ input, float* __restrict__ output,
int R, int B, int D) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * D;
if (idx >= total) return;
int b = idx / D;
int d = idx % D;
float sum = 0.0f;
for (int r = 0; r < R; r++) {
float val = input[r * B * D + b * D + d];
sum += gelu(val);
}
output[b * D + d] = sum;
}
torch::Tensor fused_gelu_sum_cuda(torch::Tensor x) {
TORCH_CHECK(x.is_cuda(), "x must be CUDA tensor");
TORCH_CHECK(x.dim() == 3, "x must be 3D");
int R = x.size(0);
int B = x.size(1);
int D = x.size(2);
auto output = torch::empty({B, D}, x.options());
int total = B * D;
int threads = 256;
int blocks = (total + threads - 1) / threads;
fused_gelu_sum_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
output.data_ptr<float>(),
R, B, D
);
return output;
}
'''
cpp_source = r'''
torch::Tensor fused_gelu_sum_cuda(torch::Tensor x);
'''
# Compile the extension
_mod = load_inline(
name="fused_gelu_sum",
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=["fused_gelu_sum_cuda"],
verbose=False,
extra_cuda_cflags=["-O3"]
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return _mod.fused_gelu_sum_cuda(x)
```