turn 0477.2s turn · 35.0s LLM · 442.2s tools · tokens in=1,091 · out=1,428 · total=3,078 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model does a virtual all-gather on the last dim: stack R tensors of shape (B, H) into (B, R*H), then Li...The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model performs a virtual all-gather on the last dimension: it stacks R tensors of shape (B, H) into (B, R*H), then applies a linear projection to (.bias).
assistant```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# CUDA kernel for fused gather + linear
cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void fused_gather_linear_kernel(
const float* __restrict__ input, // (R, B, H)
const float* __restrict__ weight, // (K, R*H)
const float* __restrict__ bias, // (K)
float* __restrict__ output, // (B, K)
int R, int B, int H, int K, int RH
) {
int b = blockIdx.x * blockDim.x + threadIdx.x;
int k = blockIdx.y * blockDim.y + threadIdx.y;
if (b >= B || k >= K) return;
float sum = bias[k];
const float* w_row = weight + k * RH;
for (int r = 0; r < R; ++r) {
const float* in_slice = input + (r * B + b) * H;
const float* w_slice = w_row + r * H;
for (int h = 0; h < H; ++h) {
sum += in_slice[h] * w_slice[h];
}
}
output[b * K + k] = sum;
}
torch::Tensor fused_gather_linear_cuda(
torch::Tensor input,
torch::Tensor weight,
torch::Tensor bias
) {
int R = input.size(0);
int B = input.size(1);
int H = input.size(2);
int K = weight.size(0);
int RH = R * H;
auto output = torch::empty({B, K}, input.options());
dim3 threads(4, 16); // 64 threads
dim3 blocks((B + 3) / 4, (K + 15) / 16);
fused_gather_linear_kernel<<<blocks, threads>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
R, B, H, K, RH
);
return output;
}
'''
cpp_source = r'''
torch::Tensor fused_gather_linear_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias);
'''
# Load the extension
fused_ext = load_inline(
name='fused_gather_linear',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['fused_gather_linear_cuda'],
verbose=False,
extra_cflags=['-O3']
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.proj = nn.Linear(4 * 32, 56, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: (R, B, H)
weight = self.proj.weight
bias = self.proj.bias
return fused_ext.fused_gather_linear_cuda(x, weight, bias)
```
The kernel above is submitted. It fuses the concatenation and linear into one kernel using direct indexing. The dimensions are tiny so loop unrolling inside is efficient, and launch config covers the (B,K)=(12,56) output space well. This avoids materializing the intermediate (B,128) tensor and associated memory traffic.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n# CUDA kernel for fused gather + linear\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void fused_gather_linear_kernel(\n const float* __restrict__ input, // (R, B, H)\n const float* __restrict__ weight, // (K, R*H)\n const float* __restrict__ bias, // (K)\n float* __restrict__ output, // (B, K)\n int R, int B, int H, int K, int RH\n) {\n int b = blockIdx.x * blockDim.x + threadIdx.x;\n int k = blockIdx.y * blockDim.y + threadIdx.y;\n \n if (b >= B || k >= K) return;\n \n float sum = bias[k];\n const float* w_row = weight + k * RH;\n \n for (int r = 0; r < R; ++r) {\n const float* in_slice = input + (r * B + b) * H;\n const float* w_slice = w_row + r * H;\n for (int h = 0; h < H; ++h) {\n sum += in_slice[h] * w_slice[h];\n }\n }\n \n output[b * K + k] = sum;\n}\n\ntorch::Tensor fused_gather_linear_cuda(\n torch::Tensor input,\n torch::Tensor weight,\n torch::Tensor bias\n) {\n int R = input.size(0);\n int B = input.size(1);\n int H = input.size(2);\n int K = weight.size(0);\n int RH = R * H;\n \n auto output = torch::empty({B, K}, input.options());\n \n dim3 threads(4, 16); // 64 threads\n dim3 blocks((B + 3) / 4, (K + 15) / 16);\n \n fused_gather_linear_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias.data_ptr<float>(),\n output.data_ptr<float>(),\n R, B, H, K, RH\n );\n \n return output;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor fused_gather_linear_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias);\n'''\n\n# Load the extension\nfused_ext = load_inline(\n name='fused_gather_linear',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['fused_gather_linear_cuda'],\n verbose=False,\n extra_cflags=['-O3']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(4 * 32, 56, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # x: (R, B, H)\n weight = self.proj.weight\n bias = self.proj.bias\n return fused_ext.fused_gather_linear_cuda(x, weight, bias)\n```"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: kernel did not compile.
Syntax error in custom generated code or ModelNew not found