incorrect
10_SampledDenseDenseMatmulEdges.py
started
2026-05-11T09:31:53Z
correctness issue:
timingkernelbench.timing
kernel mean / median / std
—
ref mean / median / std
0.05 / — / 6.30e-04 μs
roofline / occupancynsight_hardware_counters
reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Edge-only sampled dense-dense matmul.
For each graph edge (src, dst), compute a dot product between source and
destination node embeddings. This pattern appears in link prediction,
sampled attention, and graph contrastive learning.
"""
def __init__(self):
super().__init__()
def forward(
self,
src_idx: torch.Tensor,
dst_idx: torch.Tensor,
src_feat: torch.Tensor,
dst_feat: torch.Tensor,
) -> torch.Tensor:
src = src_feat[src_idx.long()]
dst = dst_feat[dst_idx.long()]
return (src * dst).sum(dim=-1)
num_src_nodes = 1024
num_dst_nodes = 768
num_edges = 16384
feat_dim = 128
def get_inputs():
p = popcorn_pri
n_edges = p.jitter_int(num_edges)
fd = p.jitter_int(feat_dim, align=8)
ns = p.jitter_int(num_src_nodes)
nd = p.jitter_int(num_dst_nodes)
src_idx = torch.randint(0, ns, (n_edges,), dtype=torch.int32)
dst_idx = torch.randint(0, nd, (n_edges,), dtype=torch.int32)
src_feat = torch.randn(ns, fd, dtype=torch.float32)
dst_feat = torch.randn(nd, fd, dtype=torch.float32)
return [src_idx, dst_idx, src_feat, dst_feat]
def get_init_inputs():
return []
turn 0584.4s turn · 15.1s LLM · 569.3s tools · tokens in=1,154 · out=620 · total=2,106 · FINAL
reasoning
The task is to replace the PyTorch model's forward with a custom CUDA kernel. The model does edge-only sampled dense-dense matmul: for each edge (src, dst), dot product between src_feat[src] and dst_f...
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 edge_dot_product_kernel(\n const int* __restrict__ src_idx,\n const int* __restrict__ dst_idx,\n const float* __restrict__ src_feat,\n const float* __restrict__ dst_feat,\n float* __restrict__ output,\n int num_edges,\n int feat_dim,\n int src_feat_stride,\n int dst_feat_stride\n) {\n int edge = blockIdx.x * blockDim.x + threadIdx.x;\n if (edge >= num_edges) return;\n\n int src = src_idx[edge];\n int dst = dst_idx[edge];\n\n const float* src_row = src_feat + src * src_feat_stride;\n const float* dst_row = dst_feat + dst * dst_feat_stride;\n\n float sum = 0.0f;\n for (int d = 0; d < feat_dim; d++) {\n sum += src_row[d] * dst_row[d];\n }\n output[edge] = sum;\n}\n\ntorch::Tensor edge_dot_product_cuda(\n torch::Tensor src_idx,\n torch::Tensor dst_idx,\n torch::Tensor src_feat,\n torch::Tensor dst_feat\n) {\n int num_edges = src_idx.size(0);\n int feat_dim = src_feat.size(1);\n int src_feat_stride = src_feat.stride(0);\n int dst_feat_stride = dst_feat.stride(0);\n\n auto output = torch::empty({num_edges}, src_feat.options());\n\n const int threads = 256;\n const int blocks = (num_edges + threads - 1) / threads;\n\n edge_dot_product_kernel<<<blocks, threads>>>(\n src_idx.data_ptr<int>(),\n dst_idx.data_ptr<int>(),\n src_feat.data_ptr<float>(),\n dst_feat.data_ptr<float>(),\n output.data_ptr<float>(),\n num_edges,\n feat_dim,\n src_feat_stride,\n dst_feat_stride\n );\n\n return output;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor edge_dot_product_cuda(\n torch::Tensor src_idx,\n torch::Tensor dst_idx,\n torch::Tensor src_feat,\n torch::Tensor dst_feat\n);\n'''\n\nmodule = load_inline(\n name='edge_dot_product',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['edge_dot_product_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(\n self,\n src_idx: torch.Tensor,\n dst_idx: torch.Tensor,\n src_feat: torch.Tensor,\n dst_feat: torch.Tensor,\n ) -> torch.Tensor:\n return self.module.edge_dot_product_cuda(src_idx, dst_idx, src_feat, dst_feat)\n"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: correctness check did not pass (? trials).
runtime_error: expected scalar type Int but found Float
final submitted kernel (99 lines, 2,439 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 edge_dot_product_kernel(
const int* __restrict__ src_idx,
const int* __restrict__ dst_idx,
const float* __restrict__ src_feat,
const float* __restrict__ dst_feat,
float* __restrict__ output,
int num_edges,
int feat_dim,
int src_feat_stride,
int dst_feat_stride
) {
int edge = blockIdx.x * blockDim.x + threadIdx.x;
if (edge >= num_edges) return;
int src = src_idx[edge];
int dst = dst_idx[edge];
const float* src_row = src_feat + src * src_feat_stride;
const float* dst_row = dst_feat + dst * dst_feat_stride;
float sum = 0.0f;
for (int d = 0; d < feat_dim; d++) {
sum += src_row[d] * dst_row[d];
}
output[edge] = sum;
}
torch::Tensor edge_dot_product_cuda(
torch::Tensor src_idx,
torch::Tensor dst_idx,
torch::Tensor src_feat,
torch::Tensor dst_feat
) {
int num_edges = src_idx.size(0);
int feat_dim = src_feat.size(1);
int src_feat_stride = src_feat.stride(0);
int dst_feat_stride = dst_feat.stride(0);
auto output = torch::empty({num_edges}, src_feat.options());
const int threads = 256;
const int blocks = (num_edges + threads - 1) / threads;
edge_dot_product_kernel<<<blocks, threads>>>(
src_idx.data_ptr<int>(),
dst_idx.data_ptr<int>(),
src_feat.data_ptr<float>(),
dst_feat.data_ptr<float>(),
output.data_ptr<float>(),
num_edges,
feat_dim,
src_feat_stride,
dst_feat_stride
);
return output;
}
'''
cpp_source = '''
torch::Tensor edge_dot_product_cuda(
torch::Tensor src_idx,
torch::Tensor dst_idx,
torch::Tensor src_feat,
torch::Tensor dst_feat
);
'''
module = load_inline(
name='edge_dot_product',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['edge_dot_product_cuda'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.module = module
def forward(
self,
src_idx: torch.Tensor,
dst_idx: torch.Tensor,
src_feat: torch.Tensor,
dst_feat: torch.Tensor,
) -> torch.Tensor:
return self.module.edge_dot_product_cuda(src_idx, dst_idx, src_feat, dst_feat)