reasoning
Let me analyze the reference implementation to understand what computation needs to be done:
1. Input: distances tensor of shape (B, N, N) where B=8, N=256
2. Output: rbf_features of shape (B, N, N, num_rbf) where num_rbf=64
Computation:
1. `d_scaled = distances.unsqueeze(-1)` - shape becomes (B, N, N, 1)
2. `envelope = 0.5 * (torch.cos(math.pi * d_scaled / self.cutoff) + 1.0)`
3. `envelope = envelope * (d_scaled < self.cutoff).float()`
4. `rbf = torch.exp(self.coeff * (d_scaled - self.offsets) ** 2)` where offsets is (num_rbf,)
5. `return rbf * envelope`
Key observations:
- The offsets tensor is broadcast against the last dimension
- Each element of distances becomes a feature vector of length num_rbf
- There are 8 * 256 * 256 = 524,288 distance values
- For each distance, we compute 64 RBF values
Memory access pattern:
- Each distance value is read once
- The offsets array (64 values) is read by every thread
- Output is 524,288 * 64 = 33,554,432 floats = 134 MB
This is a very regular, parallelizable operation. Each output element can be computed independently.
Let me write an initial kernel:
```python
import torch
import torch.nn as nn
import math
from torch.utils.cpp_extension import load_inline
# CUDA kernel
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void rbf_kernel(
const float* __restrict__ distances,
const float* __restrict__ offsets,
float coeff,
float cutoff,
float* __restrict__ output,
int total_distances,
int num_rbf
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total_distances) return;
float d = distances[idx];
// Cosine cutoff envelope
float envelope;
if (d < cutoff) {
float d_scaled = d;
envelope = 0.5f * (cosf(M_PI * d_scaled / cutoff) + 1.0f);
} else {
envelope = 0.0f;
}
// RBF expansion
for (int j = 0; j < num_rbf; j++) {
float offset = offsets[j];
float diff = d - offset;
float rbf = expf(coeff * diff * diff);
output[idx * num_rbf + j] = rbf * envelope;
}
}
torch::Tensor rbf_cuda_forward(
torch::Tensor distances,
torch::Tensor offsets,
float coeff,
float cutoff
) {
int total_distances = distances.numel();
int num_rbf = offsets.size(0);
auto output = torch::empty({total_distances, num_rbf}, distances.options());
int threads = 256;
int blocks = (total_distances + threads - 1) / threads;
rbf_kernel<<<blocks, threads>>>(
distances.data_ptr<float>(),
offsets.data_ptr<float>(),
coeff,
cutoff,
output.data_ptr<float>(),
total_distances,
num_rbf
);
return output.view(distances.sizes()).unsqueeze(-1);
}
"""
cpp_source = """
torch::Tensor rbf_cuda_forward(
torch::Tensor distances,
torch::Tensor offsets,
float coeff,
float cutoff
);
"""
rbf_extension = load_inline(
name='rbf_extension',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['rbf_cuda_forward'],
extra_cuda_cflags=['-O3', '--use_fast_math'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, num_rbf=64, cutoff=5.0):
super().__init__()
self.num_rbf = num_rbf
self.cutoff = cutoff
offsets = torch.linspace(0.0, cutoff, num_rbf)
self.register_buffer("offsets", offsets)
coeff = -0.5 / (offsets[1] - offsets[0]).item() ** 2
self.coeff = coeff
def forward(self, distances: torch.Tensor) -> torch.Tensor:
return rbf_extension.rbf_cuda_forward(distances, self.offsets, self.coeff, self.cutoff)
```
Wait, I need to think about this more carefully. The view/unsqueeze might not work directly because the output tensor has shape (total_distances, num_rbf), but we need to return shape (*distances.shape, num_rbf).
Let me think about this differently. We can allocate the output tensor directly with the right shape.
Actually, looking at the reference implementation again:
- distances shape: (B, N, N) = (8, 256, 256)
- d_scaled shape: (B, N, N, 1)
- envelope shape: (B, N, N, 1)
- offsets shape: (num_rbf,) = (64,)
- rbf shape: (B, N, N, num_rbf)
- output shape: (B, N, N, num_rbf)
So the output is indeed (*distances.shape, num_rbf) = (8, 256, 256, 64).
Let me rewrite the kernel to handle this properly:
```python
import torch
import torch.nn as nn
import math
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void rbf_kernel(
const float* __restrict__ distances,
const float* __restrict__ offsets,
float coeff,
float cutoff,
float* __restrict__ output,
int total_distances,
int num_rbf
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total_distances) return;
float d = distances[idx];
// Cosine cutoff envelope
float envelope;
if (d < cutoff) {
envelope = 0.5f * (cosf(M_PI * d / cutoff) + 1.0f);
} else {
envelope = 0.0f;
}
// RBF expansion
for (int j = 0; j < num_rbf; j++) {
float offset = offsets[j];
float diff = d - offset;
float rbf = expf(coeff * diff * diff);
output[idx * num_rbf + j] = rbf * envelope;
}
}
torch::Tensor rbf_cuda_forward(
torch::Tensor distances,
torch::Tensor offsets,
float coeff,
float cutoff
) {
auto sizes = distances.sizes().vec();
int total_distances = distances.numel();
int num_rbf = offsets.size(0);
sizes.push_back(num_rbf);
auto output = torch::empty(sizes, distances.options());
int threads = 256;
int blocks = (total_distances + threads - 1) / threads;
rbf_kernel<<<blocks, threads>>>(
distances.data_ptr<float>(),
offsets.data_ptr<float>(),
coeff,
cutoff,
output.data_ptr<float>(),
total_distances,
num_rbf
);
return output;
}
"""
cpp_source = """
torch::Tensor rbf_cuda_forward(
torch::Tensor distances,
torch::Tensor offsets,
float coeff,
float cutoff
);
"""
rbf_extension = load_inline(
name='rbf_extension',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['rbf_cuda_forward'],
extra_cuda_cflags=['-O3', '--use_fast_math'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, num_rbf=64, cutoff=5.0):
super().__init__()
self.num_rbf = num_rbf
self.cutoff = cutoff
offsets = torch.linspace(0.0, cutoff, num_rbf)
self.register_buffer("offsets", offsets)
coeff = -0.5 / (offsets[1] - offsets[0]).item() ** 2
self.coeff = coeff
def forward(self, distances: torch.Tensor) -> torch.Tensor:
return rbf_extension.rbf_cuda_forward(distances, self.offsets, self.coeff, self.cutoff)
```
This looks reasonable. Let me compile and check correctness first.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport math\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void rbf_kernel(\n const float* __restrict__ distances,\n const float* __restrict__ offsets,\n float coeff,\n float cutoff,\n float* __restrict__ output,\n int total_distances,\n int num_rbf\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= total_distances) return;\n \n float d = distances[idx];\n \n // Cosine cutoff envelope\n float envelope;\n if (d < cutoff) {\n envelope = 0.5f * (cosf(M_PI * d / cutoff) + 1.0f);\n } else {\n envelope = 0.0f;\n }\n \n // RBF expansion\n for (int j = 0; j < num_rbf; j++) {\n float offset = offsets[j];\n float diff = d - offset;\n float rbf = expf(coeff * diff * diff);\n output[idx * num_rbf + j] = rbf * envelope;\n }\n}\n\ntorch::Tensor rbf_cuda_forward(\n torch::Tensor distances,\n torch::Tensor offsets,\n float coeff,\n float cutoff\n) {\n auto sizes = distances.sizes().vec();\n int total_distances = distances.numel();\n int num_rbf = offsets.size(0);\n \n sizes.push_back(num_rbf);\n auto output = torch::empty(sizes, distances.options());\n \n int threads = 256;\n int blocks = (total_distances + threads - 1) / threads;\n \n rbf_kernel<<<blocks, threads>>>(\n distances.data_ptr<float>(),\n offsets.data_ptr<float>(),\n coeff,\n cutoff,\n output.data_ptr<float>(),\n total_distances,\n num_rbf\n );\n \n return output;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor rbf_cuda_forward(\n torch::Tensor distances,\n torch::Tensor offsets,\n float coeff,\n float cutoff\n);\n\"\"\"\n\nrbf_extension = load_inline(\n name='rbf_extension',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['rbf_cuda_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, num_rbf=64, cutoff=5.0):\n super().__init__()\n self.num_rbf = num_rbf\n self.cutoff = cutoff\n offsets = torch.linspace(0.0, cutoff, num_rbf)\n self.register_buffer(\"offsets\", offsets)\n coeff = -0.5 / (offsets[1] - offsets[0]).item() ** 2\n self.coeff = coeff\n\n def forward(self, distances: torch.Tensor) -> torch.Tensor:\n return rbf_extension.rbf_cuda_forward(distances, self.offsets, self.coeff, self.cutoff)"
}