reasoning
Now let me design a fused CUDA kernel for this CVAE. The main operations are:
1. Encoder: concat(x,y) -> Linear -> ReLU -> Linear(mu) + Linear(logvar)
2. Sampling: z = mu + std * eps, where std = exp(0.5*logvar)
3. Decoder: concat(z,y) -> Linear -> ReLU -> Linear -> x_recon
4. Loss computation: MSE loss + KL divergence
The dimensions are:
- input_dim = 256
- label_dim = 10
- hidden_dim = 128
- latent_dim = 32
- batch_size = 128
Total weights:
- enc layer: (266, 128) = 34048 params
- enc_mu: (128, 32) = 4128 params
- enc_logvar: (128, 32) = 4128 params
- dec layer 1: (42, 128) = 5376 params
- dec layer 2: (128, 256) = 33024 params
For a batch of 128, the operations are relatively small. The key optimization is to fuse all the operations into a single kernel to avoid multiple kernel launch overheads.
Let me write a kernel that:
1. Loads all weights into shared memory (they're small enough)
2. Processes the batch in parallel (one thread per output element, or one thread per sample)
3. Fuses all operations together
Actually, given the small sizes, a better approach might be:
- Use cuBLAS for the matrix multiplications (they're optimized)
- But we'd have multiple kernel launches
Or we could:
- Fuse the encoder computation into one kernel
- Fuse the decoder computation into one kernel
- Fuse the loss computation into one kernel
Let me think about the best approach...
Given the sizes:
- Batch 128, input 256, hidden 128, latent 32
- Each matrix multiplication is fairly small
A fully fused approach would:
1. For each sample in batch (or elements in parallel):
- Compute encoder hidden
- Compute mu and logvar
- Sample z
- Compute decoder output
- Accumulate loss terms
Let me write a kernel where each thread handles one sample (batch element), computing the entire forward pass for that sample, then reduce across threads for the final losses.
Wait, that might not have enough parallelism. Let me think...
128 samples * 128 hidden = 16384 operations for encoder
128 samples * 32 latent = 4096 operations for mu/logvar
128 samples * 128 hidden = 16384 operations for decoder first layer
128 samples * 256 output = 32768 operations for decoder second layer
Total ~70K FLOPs, which is tiny. The overhead of launching multiple kernels would dominate.
So the key is to fuse everything into one kernel. Given the small sizes, we could:
1. Load all weights into shared memory (they're small enough - about 80KB total)
2. Have each thread handle one sample
3. Process the forward pass for that sample
4. Reduce across threads for the loss
With 128 threads (one per sample), we'd have low occupancy. Better to have each thread handle one output element and use more threads.
Actually, let me try a different approach: use a single kernel with shared memory for weights, and have each thread compute one output element for the final result, using multiple thread blocks for better occupancy.
Hmm, this is getting complex. Let me start with a simple fused kernel and optimize from there.
I'll create a kernel where:
- Each thread block handles part of the batch
- Within a block, threads cooperate to compute the forward pass using shared memory for weights
- Final reduction for loss computation
Let me start coding:
```cpp
__global__ void cvae_forward(
const float* __restrict__ x,
const float* __restrict__ y,
const float* __restrict__ enc_weight,
const float* __restrict__ enc_bias,
const float* __restrict__ enc_mu_weight,
const float* __restrict__ enc_mu_bias,
const float* __restrict__ enc_logvar_weight,
const float* __restrict__ enc_logvar_bias,
const float* __restrict__ dec_weight,
const float* __restrict__ dec_bias,
const float* __restrict__ dec_out_weight,
const float* __restrict__ dec_out_bias,
float* __restrict__ loss_out,
float* __restrict__ recon_loss_out,
float* __restrict__ kl_loss_out,
float* __restrict__ noise, // Pre-generated random noise
int batch_size,
int input_dim,
int label_dim,
int hidden_dim,
int latent_dim
) {
// Each thread handles one batch sample
int batch_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (batch_idx >= batch_size) return;
// Load input and label
// ... compute forward pass
}
```
Actually, with the weights being small, let me try a different approach: load weights into shared memory within each block, and have each thread process multiple batch elements.
Let me write a more detailed kernel...
Actually, I realize I need to be careful about how to structure this. The key insight is:
- The model parameters are fixed (not changing during forward pass)
- We can load them once and reuse across the batch
- The random noise for reparameterization needs to be generated
For the random noise, I could:
1. Generate it in the kernel using cuRAND (but this adds complexity)
2. Generate it on the CPU and pass it in (simpler, same numerical result)
Let me go with option 2 for now, and use PyTorch to generate the random noise on the same device.
Let me write the kernel:
```cpp
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void cvae_forward_kernel(
const float* __restrict__ x,
const float* __restrict__ y,
const float* __restrict__ enc_weight, // (hidden_dim, input_dim + label_dim)
const float* __restrict__ enc_bias, // (hidden_dim,)
const float* __restrict__ enc_mu_weight, // (latent_dim, hidden_dim)
const float* __restrict__ enc_mu_bias, // (latent_dim,)
const float* __restrict__ enc_logvar_weight, // (latent_dim, hidden_dim)
const float* __restrict__ enc_logvar_bias, // (latent_dim,)
const float* __restrict__ dec_weight, // (hidden_dim, latent_dim + label_dim)
const float* __restrict__ dec_bias, // (hidden_dim,)
const float* __restrict__ dec_out_weight, // (input_dim, hidden_dim)
const float* __restrict__ dec_out_bias, // (input_dim,)
const float* __restrict__ noise, // (batch_size, latent_dim)
float* __restrict__ recon_loss_partial,
float* __restrict__ kl_loss_partial,
int batch_size,
int input_dim,
int label_dim,
int hidden_dim,
int latent_dim
) {
int batch_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (batch_idx >= batch_size) return;
int concat_dim = input_dim + label_dim;
int z_concat_dim = latent_dim + label_dim;
// Allocate local arrays for intermediate results
extern __shared__ float shared_mem[];
float* enc_hidden = shared_mem; // Not using this approach for now
// Use registers for small intermediate arrays
float hidden[128]; // hidden_dim
float mu[32]; // latent_dim
float logvar[32];
float z[32];
float dec_hidden[128];
float x_recon[256]; // input_dim
// Encoder: concat(x, y) -> Linear -> ReLU
for (int h = 0; h < hidden_dim; h++) {
float sum = enc_bias[h];
for (int i = 0; i < input_dim; i++) {
sum += x[batch_idx * input_dim + i] * enc_weight[h * concat_dim + i];
}
for (int i = 0; i < label_dim; i++) {
sum += y[batch_idx * label_dim + i] * enc_weight[h * concat_dim + input_dim + i];
}
hidden[h] = sum > 0 ? sum : 0; // ReLU
}
// Encoder mu and logvar
for (int l = 0; l < latent_dim; l++) {
float sum_mu = enc_mu_bias[l];
float sum_logvar = enc_logvar_bias[l];
for (int h = 0; h < hidden_dim; h++) {
sum_mu += hidden[h] * enc_mu_weight[l * hidden_dim + h];
sum_logvar += hidden[h] * enc_logvar_weight[l * hidden_dim + h];
}
mu[l] = sum_mu;
logvar[l] = sum_logvar;
}
// Reparameterization: z = mu + std * noise, std = exp(0.5 * logvar)
for (int l = 0; l < latent_dim; l++) {
float std = expf(0.5f * logvar[l]);
z[l] = mu[l] + std * noise[batch_idx * latent_dim + l];
}
// Decoder: concat(z, y) -> Linear -> ReLU -> Linear
for (int h = 0; h < hidden_dim; h++) {
float sum = dec_bias[h];
for (int l = 0; l < latent_dim; l++) {
sum += z[l] * dec_weight[h * z_concat_dim + l];
}
for (int l = 0; l < label_dim; l++) {
sum += y[batch_idx * label_dim + l] * dec_weight[h * z_concat_dim + latent_dim + l];
}
dec_hidden[h] = sum > 0 ? sum : 0; // ReLU
}
// Decoder output
float recon_loss = 0.0f;
for (int i = 0; i < input_dim; i++) {
float sum = dec_out_bias[i];
for (int h = 0; h < hidden_dim; h++) {
sum += dec_hidden[h] * dec_out_weight[i * hidden_dim + h];
}
x_recon[i] = sum;
float diff = sum - x[batch_idx * input_dim + i];
recon_loss += diff * diff;
}
// KL divergence: -0.5 * (1 + logvar - mu^2 - exp(logvar))
float kl_loss = 0.0f;
for (int l = 0; l < latent_dim; l++) {
kl_loss += -0.5f * (1.0f + logvar[l] - mu[l] * mu[l] - expf(logvar[l]));
}
recon_loss_partial[batch_idx] = recon_loss;
kl_loss_partial[batch_idx] = kl_loss;
}
```
This is a simple per-sample kernel. The issue is that with batch_size=128 and one thread per sample, we'd have only 128 threads total, which is very low occupancy.
Let me think about how to improve parallelism...
One approach: have each thread compute one output element for the final x_recon, then use separate kernels for different stages. But that defeats the purpose of fusion.
Another approach: use a 2D grid where one dimension is batch and the other is hidden/latent/input dimensions. This gives more threads.
Actually, let me try a different strategy:
1. For encoder hidden computation: each thread computes one hidden element for one batch sample
- Grid: (batch_size * hidden_dim) threads
2. For mu/logvar: each thread computes one latent element for one batch sample
- Grid: (batch_size * latent_dim) threads
3. etc.
But this requires synchronization between stages, which means multiple kernels or cooperative groups.
Given the small problem size, let me try a simpler approach: use a single kernel with multiple thread blocks, where each block handles a portion of the batch. Within each block, threads cooperate using shared memory.
Actually, let me just try the simple one-thread-per-sample approach first. With 128 threads per block, we can have 1 block, or with 32 threads per block, we can have 4 blocks. This might not have great occupancy, but it will avoid multiple kernel launch overheads.
Let me also consider that the kernel launch overhead might be significant compared to computation time. So fusing everything into one kernel, even with suboptimal occupancy, might still be faster than multiple optimized kernels.
Let me write the kernel code:
call → compile_kernel
callcompile_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 kernel source\ncuda_source = \"\"\"\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Fused CVAE forward kernel\n// Each thread processes one batch sample\n__global__ void cvae_forward_kernel(\n const float* __restrict__ x,\n const float* __restrict__ y,\n const float* __restrict__ enc_weight,\n const float* __restrict__ enc_bias,\n const float* __restrict__ enc_mu_weight,\n const float* __restrict__ enc_mu_bias,\n const float* __restrict__ enc_logvar_weight,\n const float* __restrict__ enc_logvar_bias,\n const float* __restrict__ dec_weight,\n const float* __restrict__ dec_bias,\n const float* __restrict__ dec_out_weight,\n const float* __restrict__ dec_out_bias,\n const float* __restrict__ noise,\n float* __restrict__ recon_loss_partial,\n float* __restrict__ kl_loss_partial,\n int batch_size,\n int input_dim,\n int label_dim,\n int hidden_dim,\n int latent_dim\n) {\n int batch_idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (batch_idx >= batch_size) return;\n \n int concat_dim = input_dim + label_dim;\n int z_concat_dim = latent_dim + label_dim;\n \n // Load input and label for this sample\n const float* x_sample = x + batch_idx * input_dim;\n const float* y_sample = y + batch_idx * label_dim;\n const float* noise_sample = noise + batch_idx * latent_dim;\n \n // Encoder: compute hidden layer\n // hidden[h] = sum_i(x[i] * W[h, i]) + sum_j(y[j] * W[h, input_dim + j]) + bias[h]\n // Then apply ReLU\n \n float recon_loss = 0.0f;\n float kl_loss = 0.0f;\n \n // We'll compute layer by layer, storing intermediate results in registers\n // For small dimensions, we can unroll loops\n \n // Compute mu and logvar directly from hidden (fused encoder)\n float mu[32]; // latent_dim max\n float logvar[32];\n \n // First compute hidden, then mu/logvar\n // To save registers, we can compute mu and logvar while computing hidden\n \n // For efficiency, let's compute hidden values on the fly for mu and logvar\n // But that would mean computing hidden twice... let's just store hidden\n \n float hidden[128]; // hidden_dim max\n \n // Compute encoder hidden layer\n for (int h = 0; h < hidden_dim; h++) {\n float sum = enc_bias[h];\n // Input contribution\n for (int i = 0; i < input_dim; i++) {\n sum += x_sample[i] * enc_weight[h * concat_dim + i];\n }\n // Label contribution\n for (int j = 0; j < label_dim; j++) {\n sum += y_sample[j] * enc_weight[h * concat_dim + input_dim + j];\n }\n hidden[h] = fmaxf(0.0f, sum); // ReLU\n }\n \n // Compute mu and logvar\n for (int l = 0; l < latent_dim; l++) {\n float sum_mu = enc_mu_bias[l];\n float sum_logvar = enc_logvar_bias[l];\n for (int h = 0; h < hidden_dim; h++) {\n sum_mu += hidden[h] * enc_mu_weight[l * hidden_dim + h];\n sum_logvar += hidden[h] * enc_logvar_weight[l * hidden_dim + h];\n }\n mu[l] = sum_mu;\n logvar[l] = sum_logvar;\n }\n \n // Reparameterization: z = mu + std * noise, std = exp(0.5 * logvar)\n float z[32];\n for (int l = 0; l < latent_dim; l++) {\n float std = expf(0.5f * logvar[l]);\n z[l] = mu[l] + std * noise_sample[l];\n }\n \n // Decoder: compute hidden layer\n float dec_hidden[128];\n for (int h = 0; h < hidden_dim; h++) {\n float sum = dec_bias[h];\n // z contribution\n for (int l = 0; l < latent_dim; l++) {\n sum += z[l] * dec_weight[h * z_concat_dim + l];\n }\n // y contribution\n for (int j = 0; j < label_dim; j++) {\n sum += y_sample[j] * dec_weight[h * z_concat_dim + latent_dim + j];\n }\n dec_hidden[h] = fmaxf(0.0f, sum); // ReLU\n }\n \n // Decoder output layer + reconstruction loss\n for (int i = 0; i < input_dim; i++) {\n float sum = dec_out_bias[i];\n for (int h = 0; h < hidden_dim; h++) {\n sum += dec_hidden[h] * dec_out_weight[i * hidden_dim + h];\n }\n float diff = sum - x_sample[i];\n recon_loss += diff * diff;\n }\n \n // KL divergence: -0.5 * sum(1 + logvar - mu^2 - exp(logvar))\n for (int l = 0; l < latent_dim; l++) {\n kl_loss += -0.5f * (1.0f + logvar[l] - mu[l] * mu[l] - expf(logvar[l]));\n }\n \n recon_loss_partial[batch_idx] = recon_loss;\n kl_loss_partial[batch_idx] = kl_loss;\n}\n\n// Reduction kernel to sum partial losses\n__global__ void reduce_loss_kernel(\n const float* __restrict__ partial,\n float* __restrict__ result,\n int n\n) {\n // Use warp shuffle for efficient reduction\n float sum = 0.0f;\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx < n) {\n sum = partial[idx];\n }\n \n // Warp-level reduction\n for (int offset = 16; offset > 0; offset >>= 1) {\n sum += __shfl_down_sync(0xffffffff, sum, offset);\n }\n \n // Only lane 0 writes result\n if ((threadIdx.x & 31) == 0) {\n atomicAdd(result, sum);\n }\n}\n\"\"\"\n\ncpp_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Forward declaration\nvoid launch_cvae_forward(\n const float* x, const float* y,\n const float* enc_weight, const float* enc_bias,\n const float* enc_mu_weight, const float* enc_mu_bias,\n const float* enc_logvar_weight, const float* enc_logvar_bias,\n const float* dec_weight, const float* dec_bias,\n const float* dec_out_weight, const float* dec_out_bias,\n const float* noise,\n float* recon_loss_out, float* kl_loss_out,\n int batch_size, int input_dim, int label_dim, int hidden_dim, int latent_dim);\n\nstd::tuple<torch::Tensor, torch::Tensor, torch::Tensor> cvae_forward_cuda(\n torch::Tensor x,\n torch::Tensor y,\n torch::Tensor enc_weight,\n torch::Tensor enc_bias,\n torch::Tensor enc_mu_weight,\n torch::Tensor enc_mu_bias,\n torch::Tensor enc_logvar_weight,\n torch::Tensor enc_logvar_bias,\n torch::Tensor dec_weight,\n torch::Tensor dec_bias,\n torch::Tensor dec_out_weight,\n torch::Tensor dec_out_bias,\n torch::Tensor noise\n) {\n int batch_size = x.size(0);\n int input_dim = x.size(1);\n int label_dim = y.size(1);\n int hidden_dim = enc_weight.size(0);\n int latent_dim = enc_mu_weight.size(0);\n \n auto options = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());\n auto recon_loss_partial = torch::zeros({batch_size}, options);\n auto kl_loss_partial = torch::zeros({batch_size}, options);\n \n launch_cvae_forward(\n x.data_ptr<float>(), y.data_ptr<float>(),\n enc_weight.data_ptr<float>(), enc_bias.data_ptr<float>(),\n enc_mu_weight.data_ptr<float>(), enc_mu_bias.data_ptr<float>(),\n enc_logvar_weight.data_ptr<float>(), enc_logvar_bias.data_ptr<float>(),\n dec_weight.data_ptr<float>(), dec_bias.data_ptr<float>(),\n dec_out_weight.data_ptr<float>(), dec_out_bias.data_ptr<float>(),\n noise.data_ptr<float>(),\n recon_loss_partial.data_ptr<float>(), kl_loss_partial.data_ptr<float>(),\n batch_size, input_dim, label_dim, hidden_dim, latent_dim\n );\n \n // Sum partial losses and divide by batch size\n float recon_loss = recon_loss_partial.sum().item<float>() / batch_size;\n float kl_loss = kl_loss_partial.sum().item<float>() / batch_size;\n \n auto loss = torch::tensor(recon_loss + kl_loss, options);\n auto recon_loss_tensor = torch::tensor(recon_loss, options);\n auto kl_loss_tensor = torch::tensor(kl_loss, options);\n \n return std::make_tuple(loss, recon_loss_tensor, kl_loss_tensor);\n}\n\"\"\"\n\n# Complete CUDA source with kernel implementation\ncuda_implementation = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void cvae_forward_kernel(\n const float* __restrict__ x,\n const float* __restrict__ y,\n const float* __restrict__ enc_weight,\n const float* __restrict__ enc_bias,\n const float* __restrict__ enc_mu_weight,\n const float* __restrict__ enc_mu_bias,\n const float* __restrict__ enc_logvar_weight,\n const float* __restrict__ enc_logvar_bias,\n const float* __restrict__ dec_weight,\n const float* __restrict__ dec_bias,\n const float* __restrict__ dec_out_weight,\n const float* __restrict__ dec_out_bias,\n const float* __restrict__ noise,\n float* __restrict__ recon_loss_partial,\n float* __restrict__ kl_loss_partial,\n int batch_size,\n int input_dim,\n int label_dim,\n int hidden_dim,\n int latent_dim\n) {\n int batch_idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (batch_idx >= batch_size) return;\n \n int concat_dim = input_dim + label_dim;\n int z_concat_dim = latent_dim + label_dim;\n \n const float* x_sample = x + batch_idx * input_dim;\n const float* y_sample = y + batch_idx * label_dim;\n const float* noise_sample = noise + batch_idx * latent_dim;\n \n float recon_loss = 0.0f;\n float kl_loss = 0.0f;\n \n float hidden[128];\n \n // Encoder hidden layer\n for (int h = 0; h < hidden_dim; h++) {\n float sum = enc_bias[h];\n for (int i = 0; i < input_dim; i++) {\n sum += x_sample[i] * enc_weight[h * concat_dim + i];\n }\n for (int j = 0; j < label_dim; j++) {\n sum += y_sample[j] * enc_weight[h * concat_dim + input_dim + j];\n }\n hidden[h] = fmaxf(0.0f, sum);\n }\n \n // mu and logvar\n float mu[32];\n float logvar[32];\n for (int l = 0; l < latent_dim; l++) {\n float sum_mu = enc_mu_bias[l];\n float sum_logvar = enc_logvar_bias[l];\n for (int h = 0; h < hidden_dim; h++) {\n sum_mu += hidden[h] * enc_mu_weight[l * hidden_dim + h];\n sum_logvar += hidden[h] * enc_logvar_weight[l * hidden_dim + h];\n }\n mu[l] = sum_mu;\n logvar[l] = sum_logvar;\n }\n \n // Reparameterization\n float z[32];\n for (int l = 0; l < latent_dim; l++) {\n float std = expf(0.5f * logvar[l]);\n z[l] = mu[l] + std * noise_sample[l];\n }\n \n // Decoder hidden layer\n float dec_hidden[128];\n for (int h = 0; h < hidden_dim; h++) {\n float sum = dec_bias[h];\n for (int l = 0; l < latent_dim; l++) {\n sum += z[l] * dec_weight[h * z_concat_dim + l];\n }\n for (int j = 0; j < label_dim; j++) {\n sum += y_sample[j] * dec_weight[h * z_concat_dim + latent_dim + j];\n }\n dec_hidden[h] = fmaxf(0.0f, sum);\n }\n \n // Decoder output + reconstruction loss\n for (int i = 0; i < input_dim; i++) {\n float sum = dec_out_bias[i];\n for (int h = 0; h < hidden_dim; h++) {\n sum += dec_hidden[h] * dec_out_weight[i * hidden_dim + h];\n }\n float diff = sum - x_sample[i];\n recon_loss += diff * diff;\n }\n \n // KL divergence\n for (int l = 0; l < latent_dim; l++) {\n kl_loss += -0.5f * (1.0f + logvar[l] - mu[l] * mu[l] - expf(logvar[l]));\n }\n \n recon_loss_partial[batch_idx] = recon_loss;\n kl_loss_partial[batch_idx] = kl_loss;\n}\n\nvoid launch_cvae_forward(\n const float* x, const float* y,\n const float* enc_weight, const float* enc_bias,\n const float* enc_mu_weight, const float* enc_mu_bias,\n const float* enc_logvar_weight, const float* enc_logvar_bias,\n const float* dec_weight, const float* dec_bias,\n const float* dec_out_weight, const float* dec_out_bias,\n const float* noise,\n float* recon_loss_partial, float* kl_loss_partial,\n int batch_size, int input_dim, int label_dim, int hidden_dim, int latent_dim\n) {\n int threads_per_block = 128;\n int num_blocks = (batch_size + threads_per_block - 1) / threads_per_block;\n \n cvae_forward_kernel<<<num_blocks, threads_per_block>>>(\n x, y, enc_weight, enc_bias, enc_mu_weight, enc_mu_bias,\n enc_logvar_weight, enc_logvar_bias, dec_weight, dec_bias,\n dec_out_weight, dec_out_bias, noise,\n recon_loss_partial, kl_loss_partial,\n batch_size, input_dim, label_dim, hidden_dim, latent_dim\n );\n}\n\nstd::tuple<torch::Tensor, torch::Tensor, torch::Tensor> cvae_forward_cuda(\n torch::Tensor x,\n torch::Tensor y,\n torch::Tensor enc_weight,\n torch::Tensor enc_bias,\n torch::Tensor enc_mu_weight,\n torch::Tensor enc_mu_bias,\n torch::Tensor enc_logvar_weight,\n torch::Tensor enc_logvar_bias,\n torch::Tensor dec_weight,\n torch::Tensor dec_bias,\n torch::Tensor dec_out_weight,\n torch::Tensor dec_out_bias,\n torch::Tensor noise\n) {\n int batch_size = x.size(0);\n int input_dim = x.size(1);\n int label_dim = y.size(1);\n int hidden_dim = enc_weight.size(0);\n int latent_dim = enc_mu_weight.size(0);\n \n auto options = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());\n auto recon_loss_partial = torch::zeros({batch_size}, options);\n auto kl_loss_partial = torch::zeros({batch_size}, options);\n \n launch_cvae_forward(\n x.data_ptr<float>(), y.data_ptr<float>(),\n enc_weight.data_ptr<float>(), enc_bias.data_ptr<float>(),\n enc_mu_weight.data_ptr<float>(), enc_mu_bias.data_ptr<float>(),\n enc_logvar_weight.data_ptr<float>(), enc_logvar_bias.data_ptr<float>(),\n dec_weight.data_ptr<float>(), dec_bias.data_ptr<float>(),\n dec_out_weight.data_ptr<float>(), dec_out_bias.data_ptr<float>(),\n noise.data_ptr<float>(),\n recon_loss_partial.data_ptr<float>(), kl_loss_partial.data_ptr<float>(),\n batch_size, input_dim, label_dim, hidden_dim, latent_dim\n );\n \n float recon_loss = recon_loss_partial.sum().item<float>() / batch_size;\n float kl_loss = kl_loss_partial.sum().item<float>() / batch_size;\n \n auto loss = torch::tensor(recon_loss + kl_loss, options);\n auto recon_loss_tensor = torch::tensor(recon_loss, options);\n auto kl_loss_tensor = torch::tensor(kl_loss, options);\n \n return std::make_tuple(loss, recon_loss_tensor, kl_loss_tensor);\n}\n\"\"\"\n\n# Load the CUDA extension\ncvae_extension = load_inline(\n name='cvae_extension',\n cpp_sources='',\n cuda_sources=cuda_implementation,\n functions=['cvae_forward_cuda'],\n extra_cuda_cflags=['-O3'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n CVAE with custom CUDA kernel for fused forward pass.\n \"\"\"\n \n def __init__(self, input_dim, label_dim, hidden_dim, latent_dim):\n super().__init__()\n self.input_dim = input_dim\n self.label_dim = label_dim\n self.hidden_dim = hidden_dim\n self.latent_dim = latent_dim\n \n # Encoder: q(z | x, y)\n self.enc = nn.Sequential(\n nn.Linear(input_dim + label_dim, hidden_dim),\n nn.ReLU(),\n )\n self.enc_mu = nn.Linear(hidden_dim, latent_dim)\n self.enc_logvar = nn.Linear(hidden_dim, latent_dim)\n \n # Decoder: p(x | z, y)\n self.dec = nn.Sequential(\n nn.Linear(latent_dim + label_dim, hidden_dim),\n nn.ReLU(),\n nn.Linear(hidden_dim, input_dim),\n )\n\n def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple:\n \"\"\"\n Args:\n x: (B, input_dim) \u2013 input data\n y: (B, label_dim) \u2013 conditioning label (one-hot)\n Returns:\n (loss, recon_loss, kl_loss)\n \"\"\"\n # Ensure inputs are contiguous and on the same device\n x = x.contiguous()\n y = y.contiguous()\n \n # Generate random noise for reparameterization\n noise = torch.randn_like(x[:, :self.latent_dim]).contiguous()\n \n # Extract weights and biases from the model\n enc_weight = self.enc[0].weight.contiguous()\n enc_bias = self.enc[0].bias.contiguous()\n enc_mu_weight = self.enc_mu.weight.contiguous()\n enc_mu_bias = self.enc_mu.bias.contiguous()\n enc_logvar_weight = self.enc_logvar.weight.contiguous()\n enc_logvar_bias = self.enc_logvar.bias.contiguous()\n dec_weight = self.dec[0].weight.contiguous()\n dec_bias = self.dec[0].bias.contiguous()\n dec_out_weight = self.dec[2].weight.contiguous()\n dec_out_bias = self.dec[2].bias.contiguous()\n \n # Call CUDA kernel\n loss, recon_loss, kl_loss = cvae_extension.cvae_forward_cuda(\n x, y,\n enc_weight, enc_bias,\n enc_mu_weight, enc_mu_bias,\n enc_logvar_weight, enc_logvar_bias,\n dec_weight, dec_bias,\n dec_out_weight, dec_out_bias,\n noise\n )\n \n return loss, recon_loss, kl_loss\n\n\n# Test dimensions\ninput_dim = 256\nlabel_dim = 10\nhidden_dim = 128\nlatent_dim = 32\nbatch_size = 128\n\n\ndef get_inputs():\n x = torch.randn(batch_size, input_dim)\n y = torch.zeros(batch_size, label_dim)\n y.scatter_(1, torch.randint(0, label_dim, (batch_size, 1)), 1.0)\n return [x, y]\n\n\ndef get_init_inputs():\n return [input_dim, label_dim, hidden_dim, latent_dim]"
}