turn 0689.8s turn · 67.4s LLM · 622.3s tools · tokens in=1,252 · out=1,704 · total=3,917 · FINAL
reasoning
The task is to replace the forward computation of this PyTorch model with a custom CUDA kernel. The model is a channel-wise gated delta attention recurrence.
I need to create a custom CUDA kernel to speed up the forward computation of this PyTorch model, which is a channel-wise gated delta attention recurrence. The original model loops over the sequence length, updating a state matrix and computing outputs for each time step. My job is to replace that with a faster CUDA kernel using Model New class, compiled with torch.utils.cpp_extension.load_inline.
The shapes involved are: for each batch and head, we have q, k, v, gate of shape [seq_len, head_dim], and state starts as [head_dim, head_dim] per batch and head, initialized to zero.
assistant```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size,
int seq_len,
int num_heads,
int head_dim
) {
int bh = blockIdx.x;
int b = bh / num_heads;
int h = bh % num_heads;
if (b >= batch_size) return;
extern __shared__ float state[];
float* s = state;
for (int i = 0; i < head_dim * head_dim; i++) {
s[i] = 0.0f;
}
__syncthreads();
const int D = head_dim;
const int BD = batch_size * num_heads * seq_len * D;
const int offset_base = (b * num_heads + h) * seq_len * D;
for (int t = 0; t < seq_len; t++) {
// Compute kv_outer and update state
for (int j = threadIdx.x; j < D * D; j += blockDim.x) {
int row = j / D;
int col = j % D;
float gt = gate[offset_base + t * D + col];
float kv = k[offset_base + t * D + row] * v[offset_base + t * D + col];
float new_state = gt * s[j] + kv;
s[j] = new_state;
}
__syncthreads();
// Compute y_t = q @ state (sum over j: q[j] * state[j, i])
for (int i = threadIdx.x; i < D; i += blockDim.x) {
float sum = 0.0f;
for (int j = 0; j < D; j++) {
sum += q[offset_base + t * D + j] * s[j * D + i];
}
int out_idx = ((b * seq_len + t) * num_heads + h) * D + i;
output[out_idx] = sum;
}
__syncthreads();
}
}
torch::Tensor gated_delta_attention_cuda(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
torch::Tensor gate
) {
int batch_size = q.size(0);
int seq_len = q.size(1);
int num_heads = q.size(2);
int head_dim = q.size(3);
auto output = torch::empty({batch_size, seq_len, num_heads, head_dim}, q.options());
int num_bh = batch_size * num_heads;
dim3 grid(num_bh);
int threads = 256;
int shared_mem = head_dim * head_dim * sizeof(float);
gated_delta_attention_kernel<<<grid, threads, shared_mem>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
gate.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
seq_len,
num_heads,
head_dim
);
return output;
}
'''
cpp_source = '''
torch::Tensor gated_delta_attention_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate);
'''
module = load_inline(
name='gated_delta_attention',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['gated_delta_attention_cuda'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.module = module
def forward(self, q, k, v, gate):
return self.module.gated_delta_attention_cuda(q, k, v, gate)
```
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 = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void gated_delta_attention_kernel(\n const float* __restrict__ q,\n const float* __restrict__ k,\n const float* __restrict__ v,\n const float* __restrict__ gate,\n float* __restrict__ output,\n int batch_size,\n int seq_len,\n int num_heads,\n int head_dim\n) {\n int bh = blockIdx.x;\n int b = bh / num_heads;\n int h = bh % num_heads;\n\n if (b >= batch_size) return;\n\n extern __shared__ float state[];\n float* s = state;\n\n for (int i = 0; i < head_dim * head_dim; i++) {\n s[i] = 0.0f;\n }\n __syncthreads();\n\n const int D = head_dim;\n const int offset_base = (b * num_heads + h) * seq_len * D;\n\n for (int t = 0; t < seq_len; t++) {\n // Compute kv_outer and update state\n for (int j = threadIdx.x; j < D * D; j += blockDim.x) {\n int row = j / D;\n int col = j % D;\n float gt = gate[offset_base + t * D + col];\n float kv = k[offset_base + t * D + row] * v[offset_base + t * D + col];\n float new_state = gt * s[j] + kv;\n s[j] = new_state;\n }\n __syncthreads();\n\n // Compute y_t = q @ state (sum over j: q[j] * state[j, i])\n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n float sum = 0.0f;\n for (int j = 0; j < D; j++) {\n sum += q[offset_base + t * D + j] * s[j * D + i];\n }\n int out_idx = ((b * seq_len + t) * num_heads + h) * D + i;\n output[out_idx] = sum;\n }\n __syncthreads();\n }\n}\n\ntorch::Tensor gated_delta_attention_cuda(\n torch::Tensor q,\n torch::Tensor k,\n torch::Tensor v,\n torch::Tensor gate\n) {\n int batch_size = q.size(0);\n int seq_len = q.size(1);\n int num_heads = q.size(2);\n int head_dim = q.size(3);\n\n auto output = torch::empty({batch_size, seq_len, num_heads, head_dim}, q.options());\n\n int num_bh = batch_size * num_heads;\n dim3 grid(num_bh);\n int threads = 256;\n int shared_mem = head_dim * head_dim * sizeof(float);\n\n gated_delta_attention_kernel<<<grid, threads, shared_mem>>>(\n q.data_ptr<float>(),\n k.data_ptr<float>(),\n v.data_ptr<float>(),\n gate.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size,\n seq_len,\n num_heads,\n head_dim\n );\n\n return output;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor gated_delta_attention_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate);\n'''\n\nmodule = load_inline(\n name='gated_delta_attention',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gated_delta_attention_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(self, q, k, v, gate):\n return self.module.gated_delta_attention_cuda(q, k, v, gate)\n```"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: kernel did not compile.
Syntax error in custom generated code or ModelNew not found