turn 0409.3s turn · 26.1s LLM · 383.2s tools · tokens in=1,279 · out=2,712 · total=3,991 · FINAL
reasoning
Let me analyze this task. I need to implement a custom CUDA kernel for RBF (Squared Exponential) covariance matrix computation.
The forward computation is:
```python
l2 = self.log_lengthscale.exp() ** 2
var = self.log_variance.exp()
# Squared distances via expansion
X1_sq = (X1 ** 2).sum(-1, keepdim=True) # (B, N, 1)
X2_sq = (X2 ** 2).sum(-1, keepdim=True) # (B, M, 1)
dist_sq = X1_sq - 2.0 * X1 @ X2.transpose(-1, -2) + X2_sq.transpose(-1, -2)
dist_sq = dist_sq.clamp(min=0.0)
return var * torch.exp(-0.5 * dist_sq / l2)
```
Input shapes:
- X1: (B, N, D) = (16, 512, 8)
- X2: (B, M, D) = (16, 128, 8)
- Output: (B, N, M) = (16, 512, 128)
The computation involves:
1. Computing squared norms of X1 and X2
2. Computing the matrix product X1 @ X2^T
3. Combining to get squared distances
4. Applying exp and scaling
For a CUDA kernel, I should:
1. Use tiled matrix multiplication approach for the main computation
2. Fuse the distance computation with the exponential
3. Optimize memory access patterns
Let me design the kernel:
Option 1: Simple approach - compute each output element K[b, i, j] in parallel
- For each thread: compute ||X1[b,i] - X2[b,j]||^2 = ||X1[b,i]||^2 - 2*dot(X1[b,i], X2[b,j]) + ||X2[b,j]||^2
- Then apply exp and scaling
I'll go with a straightforward parallelization strategy where each thread handles a single output element. This means launching a grid of threads corresponding to the output tensor dimensions. For each output position (b, i, j), I compute the squared distance by calculating the individual norms, the dot product, and combining them. Then I apply the exponential with the lengthscale and variance scaling.
The kernel computes the squared L2 distance between a point from X1 and a point from X2, applies the RBF transformation, and stores the result. I need to be careful with the indexing when loading elements from the input tensors using the batch, row, and column indices.
For each dimension, I accumulate the squared components: the norms of both points and their dot product. Then I combine these to get the squared distance, clamp it to avoid numerical issues, and apply the RBF formula with the lengthscale and variance parameters before writing to the output.
The kernel launch configuration requires a grid that covers all batches and all pairs of points between the two datasets. I'm calculating the block and grid dimensions to ensure complete coverage, then launching with the input tensors, output buffer, and hyperparameters.
For optimization, I'm considering a tiled approach to improve memory locality, but given the problem dimensions—512×128 output with 8-dimensional inputs—the simpler per-element strategy might already be efficient. The computation itself involves a few operations per element (8 fused multiply-adds for distance calculation plus an exponential), so the kernel should be relatively lightweight. The memory access pattern is the real bottleneck here—each element requires loading data from both input tensors, and these accesses aren't coalesced across threads. I should consider using shared memory to cache X2 values and improve locality.
A tiled approach would work well: partition the output matrix into tiles, preload the corresponding X1 and X2 portions into shared memory, then compute all pairwise distances within that tile. This reduces redundant global memory loads and enables coalesced access patterns.
For each batch, I'll launch enough thread blocks to cover the output matrix, where each thread computes a single output element and accumulates the dot product and squared norms across the D dimension. _sq = 0.0f;
float x2_sq = 0.0f;
float dot = 0.0f;
for (int d = 0; d < D; d++) {
float v1 = x1_ptr[d];
float v2 = x2_ptr[d];
x1_sq += v1 * v1;
x2_sq += v2 * v2;
dot += v1 * v2;
}
float dist_sq = x1_sq - 2.0f * dot + x2_sq;
if (dist_sq < 0.0f) dist_sq = 0.0f;
K_ptr[i * M + j] = variance * expf(-0.5f * dist_sq / l2);
}
}
'''
The function signature takes the input matrices, precomputed variance and lengthscale, along with batch and dimension parameters. Each thread handles one output element by iterating through all D dimensions to accumulate squared norms and the dot product, then combines these into the squared distance before applying the RBF formula. )
I should vectorize the inner loop with float4 loads for better memory throughput. If D is divisible by 4, I can process 8 elements at a time instead of iterating one-by-one, which would significantly reduce the loop overhead and improve coalescing.
Each float4 load gives me 4 consecutive dimension values, so I unpack them and accumulate the squared norms and dot products in parallel. After the vectorized section completes, any remaining dimensions get handled in a scalar cleanup loop.
Once I've computed x1_sq, x2_sq, and dot across all dimensions, I combine them into the squared distance, clamp it to prevent numerical issues, and apply the RBF formula with the lengthscale and variance parameters to get the final kernel value.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n# Compile CUDA kernel at module load time\nrbf_kernel = load_inline(\n name='rbf_kernel',\n cpp_sources='''\n#include <torch/extension.h>\n\nvoid rbf_cuda_launcher(\n const float* X1, const float* X2, float* K,\n float variance, float l2,\n int B, int N, int M, int D);\n\ntorch::Tensor rbf_forward(torch::Tensor X1, torch::Tensor X2, float variance, float l2) {\n int B = X1.size(0);\n int N = X1.size(1);\n int M = X2.size(1);\n int D = X1.size(2);\n \n auto K = torch::empty({B, N, M}, X1.options());\n \n rbf_cuda_launcher(\n X1.data_ptr<float>(),\n X2.data_ptr<float>(),\n K.data_ptr<float>(),\n variance, l2,\n B, N, M, D\n );\n \n return K;\n}\n''',\n cuda_sources='''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void rbf_kernel(\n const float* __restrict__ X1,\n const float* __restrict__ X2,\n float* __restrict__ K,\n float variance, float l2,\n int B, int N, int M, int D) {\n \n // Each thread computes one output element K[b, i, j]\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n int total = B * N * M;\n if (idx >= total) return;\n \n int j = idx % M;\n int i = (idx / M) % N;\n int b = idx / (N * M);\n \n // Pointers to input vectors\n const float* x1_ptr = X1 + (b * N + i) * D;\n const float* x2_ptr = X2 + (b * M + j) * D;\n \n float x1_sq = 0.0f;\n float x2_sq = 0.0f;\n float dot = 0.0f;\n \n // Vectorized processing using float4 (4 floats at a time)\n int d = 0;\n for (; d + 7 < D; d += 8) {\n float4 v1a = *reinterpret_cast<const float4*>(x1_ptr + d);\n float4 v2a = *reinterpret_cast<const float4*>(x2_ptr + d);\n float4 v1b = *reinterpret_cast<const float4*>(x1_ptr + d + 4);\n float4 v2b = *reinterpret_cast<const float4*>(x2_ptr + d + 4);\n \n x1_sq += v1a.x*v1a.x + v1a.y*v1a.y + v1a.z*v1a.z + v1a.w*v1a.w;\n x1_sq += v1b.x*v1b.x + v1b.y*v1b.y + v1b.z*v1b.z + v1b.w*v1b.w;\n \n x2_sq += v2a.x*v2a.x + v2a.y*v2a.y + v2a.z*v2a.z + v2a.w*v2a.w;\n x2_sq += v2b.x*v2b.x + v2b.y*v2b.y + v2b.z*v2b.z + v2b.w*v2b.w;\n \n dot += v1a.x*v2a.x + v1a.y*v2a.y + v1a.z*v2a.z + v1a.w*v2a.w;\n dot += v1b.x*v2b.x + v1b.y*v2b.y + v1b.z*v2b.z + v1b.w*v2b.w;\n }\n \n // Handle remaining elements\n for (; d < D; d++) {\n float v1 = x1_ptr[d];\n float v2 = x2_ptr[d];\n x1_sq += v1 * v1;\n x2_sq += v2 * v2;\n dot += v1 * v2;\n }\n \n // Squared distance\n float dist_sq = x1_sq - 2.0f * dot + x2_sq;\n dist_sq = fmaxf(dist_sq, 0.0f); // clamp to non-negative\n \n // RBF kernel\n K[idx] = variance * expf(-0.5f * dist_sq / l2);\n}\n\nvoid rbf_cuda_launcher(\n const float* X1, const float* X2, float* K,\n float variance, float l2,\n int B, int N, int M, int D) {\n \n int total = B * N * M;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n \n rbf_kernel<<<blocks, threads>>>(X1, X2, K, variance, l2, B, N, M, D);\n}\n''',\n functions=['rbf_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n RBF (Squared Exponential) covariance matrix computation for Gaussian\n Processes. Given two sets of input points, computes the full\n kernel matrix K[i,j] = sigma^2 * exp(-||x_i - x_j||^2 / (2 * l^2)).\n The dominant cost in GP inference and a natural target for fusion.\n \"\"\"\n\n def __init__(self, input_dim, lengthscale=1.0, variance=1.0):\n super().__init__()\n self.log_lengthscale = nn.Parameter(torch.tensor(lengthscale).log())\n self.log_variance = nn.Parameter(torch.tensor(variance).log())\n\n def forward(self, X1: torch.Tensor, X2: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n X1: (B, N, D) - first set of inputs\n X2: (B, M, D) - second set of inputs\n Returns:\n K: (B, N, M) - covariance matrix\n \"\"\"\n # Ensure contiguous float32 inputs\n X1 = X1.contiguous().float()\n X2 = X2.contiguous().float()\n \n # Compute hyperparameters\n l2 = self.log_lengthscale.exp().item() ** 2\n var = self.log_variance.exp().item()\n \n # Call fused CUDA kernel\n return rbf_kernel.rbf_forward(X1, X2, var, l2)"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (5 / 5) correctness trials all passed.
Kernel runtime: 0.07 μs
Runtime stats: mean=0.07μs median=?μs std=0.01μs
Numerical precision: max_abs_err=1.37e-06 mean_abs_err=2.84e-09 max_rel_err=7.03e-06
Memory: 37.31 MB peak (ref: 45.35 MB, ratio: 0.8227x)
Kernel launches: 9 (ref: 37, fusion ratio: 4.1111)
SOL score: 0.268