turn 0509.3s turn · 26.2s LLM · 483.1s tools · tokens in=1,256 · out=1,659 · total=3,366 · FINAL
reasoning
The task is to replace the forward computation with a custom CUDA kernel for this model that does low-rank expansion for KV activations.
The task is to create a custom CUDA kernel for a low-rank adaptation of KV activations in a transformer model. The original forward pass reshapes the hidden states, computes a latent representation via a down-projection, then projects it to key and value queries using up-projections, reshapes back, and stacks the results.
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 mla_kernel(
const float* __restrict__ hidden,
const float* __restrict__ w_down,
const float* __restrict__ w_up_k,
const float* __restrict__ w_up_v,
float* __restrict__ output,
int batch_size,
int seq_len,
int model_dim,
int rank,
int kv_heads,
int head_dim
) {
int bs = blockIdx.x;
int h = blockIdx.y;
int d = threadIdx.x;
if (d >= head_dim) return;
int B = batch_size;
int S = seq_len;
int M = model_dim;
int R = rank;
int KV = kv_heads;
int HD = head_dim;
int total_bs = B * S;
if (bs >= total_bs) return;
float k_val = 0.0f;
float v_val = 0.0f;
for (int r = 0; r < R; ++r) {
float latent = 0.0f;
for (int m = 0; m < M; ++m) {
latent += hidden[bs * M + m] * w_down[m * R + r];
}
k_val += latent * w_up_k[r * (KV * HD) + h * HD + d];
v_val += latent * w_up_v[r * (KV * HD) + h * HD + d];
}
int out_offset = bs * (2 * KV * HD) + 0 * (KV * HD) + h * HD + d;
output[out_offset] = k_val;
out_offset = bs * (2 * KV * HD) + 1 * (KV * HD) + h * HD + d;
output[out_offset] = v_val;
}
torch::Tensor mla_cuda(
torch::Tensor hidden,
torch::Tensor w_down,
torch::Tensor w_up_k,
torch::Tensor w_up_v
) {
auto batch_size = hidden.size(0);
auto seq_len = hidden.size(1);
auto model_dim = hidden.size(2);
auto rank = w_down.size(1);
auto kv_heads = 4;
auto head_dim = w_up_k.size(1) / kv_heads;
auto options = torch::TensorOptions().dtype(torch::kFloat32).device(hidden.device());
auto output = torch::empty({batch_size, seq_len, 2, kv_heads, head_dim}, options);
dim3 blocks(batch_size * seq_len, kv_heads);
dim3 threads(head_dim);
mla_kernel<<<blocks, threads>>>(
hidden.data_ptr<float>(),
w_down.data_ptr<float>(),
w_up_k.data_ptr<float>(),
w_up_v.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
seq_len,
model_dim,
rank,
kv_heads,
head_dim
);
return output;
}
'''
cpp_source = '''
torch::Tensor mla_cuda(torch::Tensor hidden, torch::Tensor w_down, torch::Tensor w_up_k, torch::Tensor w_up_v);
'''
module = load_inline(
name='mla_ext',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['mla_cuda'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.mla_cuda = module.mla_cuda
def forward(self, hidden, w_down, w_up_k, w_up_v):
return self.mla_cuda(hidden, w_down, w_up_k, w_up_v)
```
This is the final submission. I submit now.
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 mla_kernel(\n const float* __restrict__ hidden,\n const float* __restrict__ w_down,\n const float* __restrict__ w_up_k,\n const float* __restrict__ w_up_v,\n float* __restrict__ output,\n int batch_size,\n int seq_len,\n int model_dim,\n int rank,\n int kv_heads,\n int head_dim\n) {\n int bs = blockIdx.x;\n int h = blockIdx.y;\n int d = threadIdx.x;\n\n if (d >= head_dim) return;\n\n int B = batch_size;\n int S = seq_len;\n int M = model_dim;\n int R = rank;\n int KV = kv_heads;\n int HD = head_dim;\n\n int total_bs = B * S;\n if (bs >= total_bs) return;\n\n float k_val = 0.0f;\n float v_val = 0.0f;\n\n for (int r = 0; r < R; ++r) {\n float latent = 0.0f;\n for (int m = 0; m < M; ++m) {\n latent += hidden[bs * M + m] * w_down[m * R + r];\n }\n k_val += latent * w_up_k[r * (KV * HD) + h * HD + d];\n v_val += latent * w_up_v[r * (KV * HD) + h * HD + d];\n }\n\n int out_offset = bs * (2 * KV * HD) + 0 * (KV * HD) + h * HD + d;\n output[out_offset] = k_val;\n\n out_offset = bs * (2 * KV * HD) + 1 * (KV * HD) + h * HD + d;\n output[out_offset] = v_val;\n}\n\ntorch::Tensor mla_cuda(\n torch::Tensor hidden,\n torch::Tensor w_down,\n torch::Tensor w_up_k,\n torch::Tensor w_up_v\n) {\n auto batch_size = hidden.size(0);\n auto seq_len = hidden.size(1);\n auto model_dim = hidden.size(2);\n auto rank = w_down.size(1);\n auto kv_heads = 4;\n auto head_dim = w_up_k.size(1) / kv_heads;\n\n auto options = torch::TensorOptions().dtype(torch::kFloat32).device(hidden.device());\n auto output = torch::empty({batch_size, seq_len, 2, kv_heads, head_dim}, options);\n\n dim3 blocks(batch_size * seq_len, kv_heads);\n dim3 threads(head_dim);\n\n mla_kernel<<<blocks, threads>>>(\n hidden.data_ptr<float>(),\n w_down.data_ptr<float>(),\n w_up_k.data_ptr<float>(),\n w_up_v.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size,\n seq_len,\n model_dim,\n rank,\n kv_heads,\n head_dim\n );\n\n return output;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor mla_cuda(torch::Tensor hidden, torch::Tensor w_down, torch::Tensor w_up_k, torch::Tensor w_up_v);\n'''\n\nmodule = load_inline(\n name='mla_ext',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['mla_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.mla_cuda = module.mla_cuda\n\n def forward(self, hidden, w_down, w_up_k, w_up_v):\n return self.mla_cuda(hidden, w_down, w_up_k, w_up_v)\n```"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: kernel did not compile.
Syntax error in custom generated code or ModelNew not found