compile_fail
16_DeepSeekMoECombineScatter.py
started
2026-05-11T09:07:39Z
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)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Weighted expert-output combine back into token-major order.
Expert outputs are scattered back to token positions and scaled by the
routing weight associated with each expert contribution.
"""
def __init__(self):
super().__init__()
def forward(
self,
expert_hidden: torch.Tensor,
token_idx: torch.Tensor,
gates: torch.Tensor,
num_tokens: int,
) -> torch.Tensor:
out = torch.zeros(num_tokens, expert_hidden.shape[1], dtype=expert_hidden.dtype, device=expert_hidden.device)
for row in range(expert_hidden.shape[0]):
token = int(token_idx[row].item())
out[token] += gates[row] * expert_hidden[row]
return out
num_tokens = 1536
hidden_dim = 128
fanout = 2
def get_inputs():
return popcorn_pri.moe_combine_inputs(num_tokens, hidden_dim, fanout)
def get_init_inputs():
return []
turn 03773.1s turn · 3609.6s LLM · 163.5s tools · tokens in=1,148 · out=922 · total=2,406 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model combines expert outputs back into token-major order, weighted by routing gates.
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 scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
int num_rows,
int hidden_dim
) {
int row = blockIdx.x;
if (row >= num_rows) return;
int token = token_idx[row];
float gate = gates[row];
for (int d = threadIdx.x; d < hidden_dim; d += blockDim.x) {
float val = gate * expert_hidden[row * hidden_dim + d];
atomicAdd(&out[token * hidden_dim + d], val);
}
}
torch::Tensor scatter_add_cuda(
torch::Tensor expert_hidden,
torch::Tensor token_idx,
torch::Tensor gates,
int num_tokens
) {
int num_rows = expert_hidden.size(0);
int hidden_dim = expert_hidden.size(1);
auto out = torch::zeros({num_tokens, hidden_dim}, expert_hidden.options());
dim3 grid(num_rows);
dim3 block(128);
scatter_add_kernel<<<grid, block>>>(
expert_hidden.data_ptr<float>(),
token_idx.data_ptr<int>(),
gates.data_ptr<float>(),
out.data_ptr<float>(),
num_rows,
hidden_dim
);
return out;
}
'''
cpp_source = '''
torch::Tensor scatter_add_cuda(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates, int num_tokens);
'''
scatter_add = load_inline(
name='scatter_add',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['scatter_add_cuda'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, expert_hidden, token_idx, gates, num_tokens):
return scatter_add.scatter_add_cuda(expert_hidden, token_idx, gates, num_tokens)
```
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 scatter_add_kernel(\n const float* __restrict__ expert_hidden,\n const int* __restrict__ token_idx,\n const float* __restrict__ gates,\n float* __restrict__ out,\n int num_rows,\n int hidden_dim\n) {\n int row = blockIdx.x;\n if (row >= num_rows) return;\n \n int token = token_idx[row];\n float gate = gates[row];\n \n for (int d = threadIdx.x; d < hidden_dim; d += blockDim.x) {\n float val = gate * expert_hidden[row * hidden_dim + d];\n atomicAdd(&out[token * hidden_dim + d], val);\n }\n}\n\ntorch::Tensor scatter_add_cuda(\n torch::Tensor expert_hidden,\n torch::Tensor token_idx,\n torch::Tensor gates,\n int num_tokens\n) {\n int num_rows = expert_hidden.size(0);\n int hidden_dim = expert_hidden.size(1);\n \n auto out = torch::zeros({num_tokens, hidden_dim}, expert_hidden.options());\n \n dim3 grid(num_rows);\n dim3 block(128);\n \n scatter_add_kernel<<<grid, block>>>(\n expert_hidden.data_ptr<float>(),\n token_idx.data_ptr<int>(),\n gates.data_ptr<float>(),\n out.data_ptr<float>(),\n num_rows,\n hidden_dim\n );\n \n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor scatter_add_cuda(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates, int num_tokens);\n'''\n\nscatter_add = load_inline(\n name='scatter_add',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['scatter_add_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n \n def forward(self, expert_hidden, token_idx, gates, num_tokens):\n return scatter_add.scatter_add_cuda(expert_hidden, token_idx, gates, num_tokens)\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 (75 lines, 1,903 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 scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
int num_rows,
int hidden_dim
) {
int row = blockIdx.x;
if (row >= num_rows) return;
int token = token_idx[row];
float gate = gates[row];
for (int d = threadIdx.x; d < hidden_dim; d += blockDim.x) {
float val = gate * expert_hidden[row * hidden_dim + d];
atomicAdd(&out[token * hidden_dim + d], val);
}
}
torch::Tensor scatter_add_cuda(
torch::Tensor expert_hidden,
torch::Tensor token_idx,
torch::Tensor gates,
int num_tokens
) {
int num_rows = expert_hidden.size(0);
int hidden_dim = expert_hidden.size(1);
auto out = torch::zeros({num_tokens, hidden_dim}, expert_hidden.options());
dim3 grid(num_rows);
dim3 block(128);
scatter_add_kernel<<<grid, block>>>(
expert_hidden.data_ptr<float>(),
token_idx.data_ptr<int>(),
gates.data_ptr<float>(),
out.data_ptr<float>(),
num_rows,
hidden_dim
);
return out;
}
'''
cpp_source = '''
torch::Tensor scatter_add_cuda(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates, int num_tokens);
'''
scatter_add = load_inline(
name='scatter_add',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['scatter_add_cuda'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, expert_hidden, token_idx, gates, num_tokens):
return scatter_add.scatter_add_cuda(expert_hidden, token_idx, gates, num_tokens)
```