turn 0814.5s turn · 26.9s LLM · 787.6s tools · tokens in=1,131 · out=1,779 · total=3,339 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model is a three-stage MLP: h1 = relu(h0 @ W1), h2 = relu(h1 @ W2), out = h2 @ W3. With dimensions B=16...I need to create a ModelNew class that replaces the forward computation of a three-stage MLP with a custom CUDA kernel for better performance.
assistant```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
B = 16
H0 = 48
H1 = 64
H2 = 40
H3 = 24
cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void mlp_fused_kernel(
const float* __restrict__ h0,
const float* __restrict__ W1,
const float* __restrict__ W2,
const float* __restrict__ W3,
float* __restrict__ out,
int batch_size
) {
int b = blockIdx.x;
if (b >= batch_size) return;
const float* h0_row = h0 + b * 48;
// h1: 64 elements
float h1[64];
#pragma unroll
for (int j = 0; j < 64; j++) {
float sum = 0.0f;
#pragma unroll
for (int k = 0; k < 48; k++) {
sum += h0_row[k] * W1[k * 64 + j];
}
h1[j] = fmaxf(sum, 0.0f);
}
// h2: 40 elements
float h2[40];
#pragma unroll
for (int j = 0; j < 40; j++) {
float sum = 0.0f;
#pragma unroll
for (int k = 0; k < 64; k++) {
sum += h1[k] * W2[k * 40 + j];
}
h2[j] = fmaxf(sum, 0.0f);
}
// out: 24 elements
float* out_row = out + b * 24;
#pragma unroll
for (int j = 0; j < 24; j++) {
float sum = 0.0f;
#pragma unroll
for (int k = 0; k < 40; k++) {
sum += h2[k] * W3[k * 24 + j];
}
out_row[j] = sum;
}
}
torch::Tensor mlp_fused_cuda(torch::Tensor h0, torch::Tensor W1, torch::Tensor W2, torch::Tensor W3) {
int batch_size = h0.size(0);
auto options = torch::TensorOptions().dtype(torch::kFloat32).device(h0.device());
auto out = torch::empty({batch_size, 24}, options);
dim3 grid(batch_size);
dim3 block(1);
mlp_fused_kernel<<<grid, block>>>(
h0.data_ptr<float>(),
W1.data_ptr<float>(),
W2.data_ptr<float>(),
W3.data_ptr<float>(),
out.data_ptr<float>(),
batch_size
);
return out;
}
'''
cpp_source = '''
torch::Tensor mlp_fused_cuda(torch::Tensor h0, torch::Tensor W1, torch::Tensor W2, torch::Tensor W3);
'''
mlp_ext = load_inline(
name='mlp_fused',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['mlp_fused_cuda'],
verbose=False,
extra_cuda_cflags=['-O3']
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.W1 = nn.Linear(H0, H1, bias=False)
self.W2 = nn.Linear(H1, H2, bias=False)
self.W3 = nn.Linear(H2, H3, bias=False)
def forward(self, h0: torch.Tensor) -> torch.Tensor:
return mlp_ext.mlp_fused_cuda(h0, self.W1.weight, self.W2.weight, self.W3.weight)
```
The above code defines `ModelNew` with a custom fused CUDA kernel that computes the entire MLP in a single kernel launch per batch element, using register tiling for the small hidden dimensions. This eliminates intermediate tensor allocations and multiple kernel launches compared to the reference. The submission is ready.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nB = 16\nH0 = 48\nH1 = 64\nH2 = 40\nH3 = 24\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void mlp_fused_kernel(\n const float* __restrict__ h0,\n const float* __restrict__ W1,\n const float* __restrict__ W2,\n const float* __restrict__ W3,\n float* __restrict__ out,\n int batch_size\n) {\n int b = blockIdx.x;\n if (b >= batch_size) return;\n\n const float* h0_row = h0 + b * 48;\n \n // h1: 64 elements\n float h1[64];\n #pragma unroll\n for (int j = 0; j < 64; j++) {\n float sum = 0.0f;\n #pragma unroll\n for (int k = 0; k < 48; k++) {\n sum += h0_row[k] * W1[k * 64 + j];\n }\n h1[j] = fmaxf(sum, 0.0f);\n }\n \n // h2: 40 elements\n float h2[40];\n #pragma unroll\n for (int j = 0; j < 40; j++) {\n float sum = 0.0f;\n #pragma unroll\n for (int k = 0; k < 64; k++) {\n sum += h1[k] * W2[k * 40 + j];\n }\n h2[j] = fmaxf(sum, 0.0f);\n }\n \n // out: 24 elements\n float* out_row = out + b * 24;\n #pragma unroll\n for (int j = 0; j < 24; j++) {\n float sum = 0.0f;\n #pragma unroll\n for (int k = 0; k < 40; k++) {\n sum += h2[k] * W3[k * 24 + j];\n }\n out_row[j] = sum;\n }\n}\n\ntorch::Tensor mlp_fused_cuda(torch::Tensor h0, torch::Tensor W1, torch::Tensor W2, torch::Tensor W3) {\n int batch_size = h0.size(0);\n auto options = torch::TensorOptions().dtype(torch::kFloat32).device(h0.device());\n auto out = torch::empty({batch_size, 24}, options);\n \n dim3 grid(batch_size);\n dim3 block(1);\n \n mlp_fused_kernel<<<grid, block>>>(\n h0.data_ptr<float>(),\n W1.data_ptr<float>(),\n W2.data_ptr<float>(),\n W3.data_ptr<float>(),\n out.data_ptr<float>(),\n batch_size\n );\n \n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor mlp_fused_cuda(torch::Tensor h0, torch::Tensor W1, torch::Tensor W2, torch::Tensor W3);\n'''\n\nmlp_ext = load_inline(\n name='mlp_fused',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['mlp_fused_cuda'],\n verbose=False,\n extra_cuda_cflags=['-O3']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.W1 = nn.Linear(H0, H1, bias=False)\n self.W2 = nn.Linear(H1, H2, bias=False)\n self.W3 = nn.Linear(H2, H3, bias=False)\n\n def forward(self, h0: torch.Tensor) -> torch.Tensor:\n return mlp_ext.mlp_fused_cuda(h0, self.W1.weight, self.W2.weight, self.W3.weight)\n```"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: kernel did not compile.
Syntax error in custom generated code or ModelNew not found