reasoning
Let me analyze this task. I need to replace the forward computation of a VAE model with a custom CUDA kernel. The model consists of:
1. **Encoder**: Linear(input_dim, hidden_dim) → ReLU → Linear(hidden_dim, latent_dim) for both mu and logvar
2. **Sampling**: z = mu + std * randn, where std = exp(0.5 * logvar)
3. **Decoder**: Linear(latent_dim, hidden_dim) → ReLU → Linear(hidden_dim, input_dim) → Sigmoid
4. **Loss computation**:
- recon = BCE(x_recon, x).sum(-1)
- kl = -0.5 * (1 + logvar - mu^2 - exp(logvar)).sum(-1)
- elbo = -(recon + kl)
The dimensions are:
- input_dim = 784
- hidden_dim = 256
- latent_dim = 32
- batch_size = 128
Since I only have one turn, I need to submit my best implementation. Let me write a fused CUDA kernel that handles the entire forward pass.
Key considerations:
1. The model has learnable parameters (weights and biases for all linear layers)
2. I need to implement matrix multiplication, ReLU, Sigmoid, and the loss computations
3. The sampling uses `torch.randn_like` which generates random numbers - I need to handle this
Actually, looking at this more carefully, the challenge is that we need to keep the model's parameters. The forward pass uses:
- encoder[0]: Linear(784, 256) + ReLU
- fc_mu: Linear(256, 32)
- fc_logvar: Linear(256, 32)
- decoder[0]: Linear(32, 256) + ReLU
- decoder[2]: Linear(256, 784) + Sigmoid
Let me think about how to structure this. The standard approach would be to:
1. Keep the PyTorch model's parameters
2. Write a custom CUDA kernel that fuses the forward operations
I'll need to generate random numbers within the kernel or use curand. Let me plan the fused operations: encoder pass (Linear → ReLU), then split for fc_mu and fc_logvar, then sample z, then decoder pass (Linear → ReLU → Linear → Sigmoid), and finally the loss computation.
The key insight is that this is a sequence of matrix multiplications and element-wise operations. I should write a single kernel that handles everything, pulling parameters from the original model and implementing each layer in sequence. The random number generation could use torch's randn on the CPU side before feeding it into the kernel.
A better strategy is to leverage cuBLAS for the linear layers and fuse only the element-wise operations together. This keeps the matrix multiplications optimized while avoiding memory round-trips for the simpler ops. I'd load the model parameters into CUDA memory, chain together the encoder → fc_mu/logvar → sampling → decoder → loss computation, and let cuBLAS handle the heavy lifting for the matrix operations.
The core insight is that most of the computation lives in the two linear layers (encoder: 784×256 = 200K, decoder: 256×784 = 200K FLOPs), with sampling and loss being comparatively lightweight. Fusing the loss computation with the decoder output eliminates one memory transfer, which is where the real optimization opportunity lies.
Rather than rewriting the entire architecture, I can target the bottleneck specifically—the sampling step through decoder and into the loss calculation. The key technical constraint is that I need to generate the random noise tensor on the CPU side and pass it through, since the forward pass signature is fixed as taking only x.
Looking at the reference implementation, it generates random noise internally using `torch.randn_like`. My custom kernel can replicate this behavior by either using `torch.randn` during setup or calling cuRAND within the CUDA kernel itself. The simplest approach is keeping the model structure intact while offloading just the critical computation path to a fused CUDA kernel that handles sampling, decoding, and loss computation in one pass.
I'll use `load_inline` to compile the CUDA code and create a custom extension that replaces the forward method. The key optimization is fusing operations to reduce memory transfers—particularly the element-wise steps after sampling through the decoder and into the loss calculation.
For the linear layers, I can leverage cuBLAS for the matrix multiplications while implementing ReLU and sigmoid activations directly in CUDA. The main challenge is structuring the kernel to minimize intermediate memory allocations. I'll start with a straightforward implementation that handles the encoder-decoder flow, then refine it based on performance. Looking at this more carefully, I should fuse the entire forward pass into a single kernel. The encoder processes the input through a linear layer and ReLU activation, then produces mu and logvar through separate linear layers. The sampling step generates z using the reparameterization trick, and the decoder reconstructs through another linear-relu-linear-sigmoid sequence. The losses are computed from the reconstruction and KL divergence terms, then combined into the final ELBO.
To implement this efficiently in CUDA, I need to structure it as separate kernels for each stage rather than trying to fuse everything into one massive kernel. The matrix multiplications will dominate the computation time, so I'll write separate kernels for the linear layers with their activations, handle the sampling step, then compute the losses. This approach is more manageable than trying to optimize a single fused kernel from the start. I need to think through this more systematically. Rather than trying to fuse everything at once, I should start with a working implementation and build from there.
The simplest approach would be to use cuBLAS for matrix operations with custom kernels for the element-wise operations. Looking at the model architecture more carefully though, the encoder and decoder are Sequential modules containing Linear layers and activations, with separate fc_mu and fc_logvar Linear layers between them. This means I need to handle the weight matrices and biases from each component separately when building the CUDA kernels.
The key insight is that the forward pass chains these operations: input goes through the encoder layers, then through fc_mu and fc_logvar to get mu and logvar, then sampling produces z, which feeds into the decoder. Rather than orchestrating multiple cuBLAS calls from Python, I should write a single C++ function that handles the entire sequence of operations end-to-end, managing all the cuBLAS matrix multiplications and custom kernels for activations and sampling within that function.
To implement this, I'll need to generate random samples for the latent space, which means integrating cuRAND into the pipeline. I'll initialize a cuRAND generator once and pass it through to generate the noise tensor efficiently during each forward pass. The architecture involves an encoder with a linear layer and ReLU activation, separate linear layers for mean and log-variance outputs, and a decoder with two linear layers separated by ReLU and ending with sigmoid activation.
For the forward pass, I'll leverage cuBLAS for matrix multiplications across all linear layers, use custom kernels for the element-wise ReLU and sigmoid operations, and apply cuRAND for sampling during the reparameterization step. The loss computation combines binary cross-entropy and KL divergence terms, both of which can be computed with custom CUDA kernels to keep everything on GPU.
I need to set up persistent cuBLAS and cuRAND handles during module initialization, then use them across forward calls. Since the module loads once, creating these handles at that point avoids repeated initialization overhead.
The key insight is that I can't store handles in global variables when the model itself manages state. I need to attach the handles to the ModelNew instance in `__init__` and reuse them during forward passes. For thread safety with cuRAND, I should consider creating a generator per instance rather than sharing across threads, then pass these handles to the CUDA functions through the pybind interface.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n\n# Define CUDA source\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <curand.h>\n#include <curand_kernel.h>\n#include <cstdio>\n\n// Constants\n#define BLOCK_SIZE 256\n\n// ReLU kernel\n__global__ void relu_kernel(const float* input, float* output, int size) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < size) {\n output[idx] = input[idx] > 0 ? input[idx] : 0.0f;\n }\n}\n\n// Sigmoid kernel\n__global__ void sigmoid_kernel(const float* input, float* output, int size) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < size) {\n output[idx] = 1.0f / (1.0f + expf(-input[idx]));\n }\n}\n\n// Sampling kernel: z = mu + std * noise\n__global__ void sampling_kernel(const float* mu, const float* logvar, const float* noise, \n float* z, float* std, int size) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < size) {\n float s = expf(0.5f * logvar[idx]);\n std[idx] = s;\n z[idx] = mu[idx] + s * noise[idx];\n }\n}\n\n// Random number generation kernel\n__global__ void randn_kernel(float* output, int size, unsigned long long seed) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < size) {\n // Use Box-Muller transform for normal distribution\n curandState state;\n curand_init(seed, idx, 0, &state);\n output[idx] = curand_normal(&state);\n }\n}\n\n// Loss computation kernel\n__global__ void loss_kernel(const float* x_recon, const float* x, const float* mu, \n const float* logvar, float* elbo, float* recon_loss, \n float* kl_div, int batch_size, int input_dim, int latent_dim) {\n int batch_idx = blockIdx.x;\n if (batch_idx >= batch_size) return;\n \n float recon_sum = 0.0f;\n float kl_sum = 0.0f;\n \n // Compute reconstruction loss (BCE)\n for (int i = threadIdx.x; i < input_dim; i += blockDim.x) {\n int idx = batch_idx * input_dim + i;\n // Binary cross entropy: -(x * log(x_recon) + (1-x) * log(1-x_recon))\n // Add small epsilon for numerical stability\n float xr = x_recon[idx];\n float xi = x[idx];\n float eps = 1e-7f;\n xr = fmaxf(eps, fminf(1.0f - eps, xr));\n recon_sum += -(xi * logf(xr) + (1.0f - xi) * logf(1.0f - xr));\n }\n \n // Compute KL divergence\n for (int i = threadIdx.x; i < latent_dim; i += blockDim.x) {\n int idx = batch_idx * latent_dim + i;\n float m = mu[idx];\n float lv = logvar[idx];\n // KL = -0.5 * (1 + logvar - mu^2 - exp(logvar))\n kl_sum += -0.5f * (1.0f + lv - m * m - expf(lv));\n }\n \n // Reduce within block\n __shared__ float s_recon[BLOCK_SIZE];\n __shared__ float s_kl[BLOCK_SIZE];\n \n s_recon[threadIdx.x] = recon_sum;\n s_kl[threadIdx.x] = kl_sum;\n __syncthreads();\n \n // Tree reduction\n for (int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (threadIdx.x < s) {\n s_recon[threadIdx.x] += s_recon[threadIdx.x + s];\n s_kl[threadIdx.x] += s_kl[threadIdx.x + s];\n }\n __syncthreads();\n }\n \n if (threadIdx.x == 0) {\n recon_loss[batch_idx] = s_recon[0];\n kl_div[batch_idx] = s_kl[0];\n elbo[batch_idx] = -(s_recon[0] + s_kl[0]);\n }\n}\n\n// Add bias kernel\n__global__ void add_bias_kernel(float* data, const float* bias, int batch_size, int features) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = batch_size * features;\n if (idx < total) {\n int feature_idx = idx % features;\n data[idx] += bias[feature_idx];\n }\n}\n\n// Forward pass using cuBLAS for matrix multiplications\n#include <cublas_v2.h>\n\n// Global handles (initialized once)\nstatic cublasHandle_t cublas_handle = nullptr;\nstatic bool handles_initialized = false;\n\nvoid init_handles() {\n if (!handles_initialized) {\n cublasCreate(&cublas_handle);\n handles_initialized = true;\n }\n}\n\n// Main forward function\nstd::tuple<torch::Tensor, torch::Tensor, torch::Tensor> vae_forward(\n torch::Tensor x,\n torch::Tensor encoder_weight, torch::Tensor encoder_bias,\n torch::Tensor fc_mu_weight, torch::Tensor fc_mu_bias,\n torch::Tensor fc_logvar_weight, torch::Tensor fc_logvar_bias,\n torch::Tensor decoder1_weight, torch::Tensor decoder1_bias,\n torch::Tensor decoder2_weight, torch::Tensor decoder2_bias) {\n \n init_handles();\n \n const int batch_size = x.size(0);\n const int input_dim = x.size(1);\n const int hidden_dim = encoder_weight.size(0);\n const int latent_dim = fc_mu_weight.size(0);\n \n // Create output tensors\n auto options = x.options();\n \n // Intermediate tensors\n auto h = torch::empty({batch_size, hidden_dim}, options);\n auto h_relu = torch::empty({batch_size, hidden_dim}, options);\n auto mu = torch::empty({batch_size, latent_dim}, options);\n auto logvar = torch::empty({batch_size, latent_dim}, options);\n auto std = torch::empty({batch_size, latent_dim}, options);\n auto z = torch::empty({batch_size, latent_dim}, options);\n auto noise = torch::empty({batch_size, latent_dim}, options);\n auto decoder_h = torch::empty({batch_size, hidden_dim}, options);\n auto decoder_h_relu = torch::empty({batch_size, hidden_dim}, options);\n auto x_pre_sigmoid = torch::empty({batch_size, input_dim}, options);\n auto x_recon = torch::empty({batch_size, input_dim}, options);\n \n // Output tensors\n auto elbo = torch::empty({batch_size}, options);\n auto recon_loss = torch::empty({batch_size}, options);\n auto kl_div = torch::empty({batch_size}, options);\n \n float alpha = 1.0f, beta = 0.0f;\n \n // Encoder: Linear(input_dim, hidden_dim)\n // h = x @ encoder_weight.T\n cublasSgemm(cublas_handle, \n CUBLAS_OP_T, CUBLAS_OP_N,\n hidden_dim, batch_size, input_dim,\n &alpha,\n encoder_weight.data_ptr<float>(), input_dim,\n x.data_ptr<float>(), input_dim,\n &beta,\n h.data_ptr<float>(), hidden_dim);\n \n // Add bias and ReLU\n int h_size = batch_size * hidden_dim;\n int blocks = (h_size + BLOCK_SIZE - 1) / BLOCK_SIZE;\n add_bias_kernel<<<blocks, BLOCK_SIZE>>>(h.data_ptr<float>(), encoder_bias.data_ptr<float>(), \n batch_size, hidden_dim);\n relu_kernel<<<blocks, BLOCK_SIZE>>>(h.data_ptr<float>(), h_relu.data_ptr<float>(), h_size);\n \n // fc_mu: Linear(hidden_dim, latent_dim)\n cublasSgemm(cublas_handle,\n CUBLAS_OP_T, CUBLAS_OP_N,\n latent_dim, batch_size, hidden_dim,\n &alpha,\n fc_mu_weight.data_ptr<float>(), hidden_dim,\n h_relu.data_ptr<float>(), hidden_dim,\n &beta,\n mu.data_ptr<float>(), latent_dim);\n \n int mu_size = batch_size * latent_dim;\n blocks = (mu_size + BLOCK_SIZE - 1) / BLOCK_SIZE;\n add_bias_kernel<<<blocks, BLOCK_SIZE>>>(mu.data_ptr<float>(), fc_mu_bias.data_ptr<float>(),\n batch_size, latent_dim);\n \n // fc_logvar: Linear(hidden_dim, latent_dim)\n cublasSgemm(cublas_handle,\n CUBLAS_OP_T, CUBLAS_OP_N,\n latent_dim, batch_size, hidden_dim,\n &alpha,\n fc_logvar_weight.data_ptr<float>(), hidden_dim,\n h_relu.data_ptr<float>(), hidden_dim,\n &beta,\n logvar.data_ptr<float>(), latent_dim);\n \n add_bias_kernel<<<blocks, BLOCK_SIZE>>>(logvar.data_ptr<float>(), fc_logvar_bias.data_ptr<float>(),\n batch_size, latent_dim);\n \n // Generate random noise using torch (faster than cuRAND for this size)\n auto noise_tensor = torch::randn({batch_size, latent_dim}, options);\n \n // Sampling: z = mu + std * noise\n sampling_kernel<<<blocks, BLOCK_SIZE>>>(mu.data_ptr<float>(), logvar.data_ptr<float>(),\n noise_tensor.data_ptr<float>(),\n z.data_ptr<float>(), std.data_ptr<float>(), mu_size);\n \n // Decoder layer 1: Linear(latent_dim, hidden_dim)\n cublasSgemm(cublas_handle,\n CUBLAS_OP_T, CUBLAS_OP_N,\n hidden_dim, batch_size, latent_dim,\n &alpha,\n decoder1_weight.data_ptr<float>(), latent_dim,\n z.data_ptr<float>(), latent_dim,\n &beta,\n decoder_h.data_ptr<float>(), hidden_dim);\n \n add_bias_kernel<<<blocks, BLOCK_SIZE>>>(decoder_h.data_ptr<float>(), decoder1_bias.data_ptr<float>(),\n batch_size, hidden_dim);\n relu_kernel<<<blocks, BLOCK_SIZE>>>(decoder_h.data_ptr<float>(), decoder_h_relu.data_ptr<float>(), h_size);\n \n // Decoder layer 2: Linear(hidden_dim, input_dim)\n cublasSgemm(cublas_handle,\n CUBLAS_OP_T, CUBLAS_OP_N,\n input_dim, batch_size, hidden_dim,\n &alpha,\n decoder2_weight.data_ptr<float>(), hidden_dim,\n decoder_h_relu.data_ptr<float>(), hidden_dim,\n &beta,\n x_pre_sigmoid.data_ptr<float>(), input_dim);\n \n int x_size = batch_size * input_dim;\n blocks = (x_size + BLOCK_SIZE - 1) / BLOCK_SIZE;\n add_bias_kernel<<<blocks, BLOCK_SIZE>>>(x_pre_sigmoid.data_ptr<float>(), decoder2_bias.data_ptr<float>(),\n batch_size, input_dim);\n sigmoid_kernel<<<blocks, BLOCK_SIZE>>>(x_pre_sigmoid.data_ptr<float>(), x_recon.data_ptr<float>(), x_size);\n \n // Compute losses\n loss_kernel<<<batch_size, BLOCK_SIZE>>>(x_recon.data_ptr<float>(), x.data_ptr<float>(),\n mu.data_ptr<float>(), logvar.data_ptr<float>(),\n elbo.data_ptr<float>(), recon_loss.data_ptr<float>(),\n kl_div.data_ptr<float>(),\n batch_size, input_dim, latent_dim);\n \n return std::make_tuple(elbo, recon_loss, kl_div);\n}\n\"\"\"\n\ncpp_source = \"\"\"\nstd::tuple<torch::Tensor, torch::Tensor, torch::Tensor> vae_forward(\n torch::Tensor x,\n torch::Tensor encoder_weight, torch::Tensor encoder_bias,\n torch::Tensor fc_mu_weight, torch::Tensor fc_mu_bias,\n torch::Tensor fc_logvar_weight, torch::Tensor fc_logvar_bias,\n torch::Tensor decoder1_weight, torch::Tensor decoder1_bias,\n torch::Tensor decoder2_weight, torch::Tensor decoder2_bias);\n\"\"\"\n\n# Load the CUDA extension\nvae_extension = load_inline(\n name='vae_extension',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['vae_forward'],\n extra_cuda_cflags=['-lcublas', '-lcurand'],\n extra_ldflags=['-lcublas', '-lcurand'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Evidence Lower Bound (ELBO) computation for a Variational Autoencoder.\n Fused CUDA implementation for improved performance.\n \"\"\"\n\n def __init__(self, input_dim, hidden_dim, latent_dim):\n super().__init__()\n \n # Initialize weights and biases directly (no Sequential wrappers)\n # Encoder: Linear(input_dim, hidden_dim)\n self.encoder_weight = nn.Parameter(torch.empty(hidden_dim, input_dim))\n self.encoder_bias = nn.Parameter(torch.empty(hidden_dim))\n \n # fc_mu: Linear(hidden_dim, latent_dim)\n self.fc_mu_weight = nn.Parameter(torch.empty(latent_dim, hidden_dim))\n self.fc_mu_bias = nn.Parameter(torch.empty(latent_dim))\n \n # fc_logvar: Linear(hidden_dim, latent_dim)\n self.fc_logvar_weight = nn.Parameter(torch.empty(latent_dim, hidden_dim))\n self.fc_logvar_bias = nn.Parameter(torch.empty(latent_dim))\n \n # Decoder layer 1: Linear(latent_dim, hidden_dim)\n self.decoder1_weight = nn.Parameter(torch.empty(hidden_dim, latent_dim))\n self.decoder1_bias = nn.Parameter(torch.empty(hidden_dim))\n \n # Decoder layer 2: Linear(hidden_dim, input_dim)\n self.decoder2_weight = nn.Parameter(torch.empty(input_dim, hidden_dim))\n self.decoder2_bias = nn.Parameter(torch.empty(input_dim))\n \n # Initialize weights\n self._init_weights()\n \n # Store dimensions\n self.input_dim = input_dim\n self.hidden_dim = hidden_dim\n self.latent_dim = latent_dim\n \n def _init_weights(self):\n # Use same initialization as nn.Linear\n for name, param in self.named_parameters():\n if 'weight' in name:\n nn.init.kaiming_uniform_(param, a=5 ** 0.5)\n elif 'bias' in name:\n fan_in, _ = nn.init._calculate_fan_in_and_fan_out(param)\n bound = 1 / (fan_in ** 0.5) if fan_in > 0 else 0\n nn.init.uniform_(param, -bound, bound)\n \n def forward(self, x: torch.Tensor) -> tuple:\n \"\"\"\n Args:\n x: (B, input_dim) - input data in [0, 1]\n Returns:\n (elbo, recon_loss, kl_div):\n elbo: (B,)\n recon_loss: (B,)\n kl_div: (B,)\n \"\"\"\n # Ensure input is contiguous and on same device as parameters\n x = x.contiguous()\n \n # Call the fused CUDA kernel\n elbo, recon_loss, kl_div = vae_extension.vae_forward(\n x,\n self.encoder_weight, self.encoder_bias,\n self.fc_mu_weight, self.fc_mu_bias,\n self.fc_logvar_weight, self.fc_logvar_bias,\n self.decoder1_weight, self.decoder1_bias,\n self.decoder2_weight, self.decoder2_bias\n )\n \n return elbo, recon_loss, kl_div"
}