reasoning
Now let me understand the reference implementation and write a fused CUDA kernel.
The operations are:
1. LayerNorm on (channels, seq_len) - treating the last two dimensions
2. GELU activation
3. Dilated Conv1d with kernel_size=3, dilation=4
4. Residual addition
For the fused kernel, I need to consider:
- LayerNorm requires computing mean and variance across (channels, seq_len) for each batch element
- GELU is element-wise
- Conv1d is the expensive part - dilated convolution
- Residual addition is element-wise
Let me think about memory requirements:
- Input: [32, 256, 4096] = 32M elements = 128MB in fp32
- Output: same size
For Conv1d with dilation=4, kernel_size=3, the effective receptive field is 1 + 4*(3-1) = 9.
With padding=4, the output has the same size as input.
Let me think about how to fuse these operations efficiently:
Option 1: Fuse everything into one kernel
- Challenge: Conv1d requires different threads to access overlapping input regions
- The norm+GELU can be fused easily, but conv is trickier
Option 2: Separate kernels
- Kernel 1: LayerNorm + GELU
- Kernel 2: Dilated Conv1d + residual add
Let me start with a simpler approach - fusing LayerNorm + GELU into one kernel, and doing Conv1d + residual in another. This should already give speedup.
Actually, let me think more carefully. The Conv1d is channels->channels, so it's a depthwise-like operation if channels are preserved. Let me check - yes, in_channels=channels, out_channels=channels.
A 1x1 conv would be cheap, but kernel_size=3 with dilation=4 means each output element needs 3 input elements from the same channel but different sequence positions.
Let me write a fused kernel that:
1. Computes LayerNorm per-batch
2. Applies GELU
3. Computes dilated Conv1d
4. Adds residual
For the fused approach, I'll need to load the input twice - once for normalization statistics, and once for the actual computation. Or I can compute statistics first, then launch a second kernel that does the fused norm+gelu+conv+add.
Actually, the best approach might be:
1. One kernel to compute per-batch mean and variance (small output)
2. One fused kernel that does: normalize, gelu, conv, add residual
This way we only load the input twice (once for stats, once for computation), which is optimal.
Let me write this:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
import math
# Define the CUDA kernel
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
// GELU approximation
__device__ __forceinline__ float gelu(float x) {
// GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
constexpr float SQRT_2_PI = 0.7978845608f;
float x3 = x * x * x;
float inner = SQRT_2_PI * (x + 0.044715f * x3);
return 0.5f * x * (1.0f + tanhf(inner));
}
// Kernel to compute mean and variance per batch element
__global__ void compute_stats_kernel(
const float* __restrict__ input,
float* __restrict__ mean,
float* __restrict__ rstd,
int batch_size,
int channels,
int seq_len
) {
int batch_idx = blockIdx.x;
if (batch_idx >= batch_size) return;
int tid = threadIdx.x;
int numel = channels * seq_len;
// Use shared memory for reduction
extern __shared__ float sdata[];
float* s_sum = sdata;
float* s_sum_sq = sdata + blockDim.x;
float local_sum = 0.0f;
float local_sum_sq = 0.0f;
// Stride loop over elements
for (int i = tid; i < numel; i += blockDim.x) {
int c = i / seq_len;
int s = i % seq_len;
int idx = batch_idx * channels * seq_len + c * seq_len + s;
float val = input[idx];
local_sum += val;
local_sum_sq += val * val;
}
s_sum[tid] = local_sum;
s_sum_sq[tid] = local_sum_sq;
__syncthreads();
// Block reduction
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) {
s_sum[tid] += s_sum[tid + s];
s_sum_sq[tid] += s_sum_sq[tid + s];
}
__syncthreads();
}
if (tid == 0) {
float mean_val = s_sum[0] / numel;
float var_val = s_sum_sq[0] / numel - mean_val * mean_val;
mean[batch_idx] = mean_val;
rstd[batch_idx] = 1.0f / sqrtf(var_val + 1e-5f);
}
}
// Fused kernel: normalize, gelu, conv, add residual
__global__ void fused_norm_gelu_conv_add_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
const float* __restrict__ mean,
const float* __restrict__ rstd,
float* __restrict__ output,
int batch_size,
int channels,
int seq_len,
int kernel_size,
int dilation,
int padding
) {
int batch_idx = blockIdx.z;
int channel_idx = blockIdx.y;
int seq_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (batch_idx >= batch_size || channel_idx >= channels || seq_idx >= seq_len) return;
float mean_val = mean[batch_idx];
float rstd_val = rstd[batch_idx];
// Compute normalized + gelu values needed for conv
// For kernel_size=3, dilation=4, we need positions: seq_idx-4, seq_idx, seq_idx+4
// But we need to handle padding
int w_idx = channel_idx * kernel_size; // weight index for this channel
float conv_sum = 0.0f;
for (int k = 0; k < kernel_size; k++) {
int input_seq = seq_idx - padding + k * dilation;
if (input_seq >= 0 && input_seq < seq_len) {
int idx = batch_idx * channels * seq_len + channel_idx * seq_len + input_seq;
float val = input[idx];
float norm_val = (val - mean_val) * rstd_val;
float gelu_val = gelu(norm_val);
conv_sum += gelu_val * weight[w_idx + k];
}
}
float out_val = conv_sum + bias[channel_idx];
// Add residual
int residual_idx = batch_idx * channels * seq_len + channel_idx * seq_len + seq_idx;
out_val += input[residual_idx];
output[residual_idx] = out_val;
}
"""
cpp_source = """
torch::Tensor compute_stats(torch::Tensor input);
torch::Tensor fused_norm_gelu_conv_add(
torch::Tensor input,
torch::Tensor weight,
torch::Tensor bias,
torch::Tensor mean,
torch::Tensor rstd,
int kernel_size,
int dilation,
int padding
);
"""
functions = ["compute_stats", "fused_norm_gelu_conv_add"]
# I need to implement the host functions...
```
Wait, I need to be more careful. The cpp_source should define the C++ host functions that call the CUDA kernels. Let me restructure this properly.
Actually, let me look at the torch.utils.cpp_extension.load_inline documentation. The cuda_sources and cpp_sources are combined, and functions lists the functions to bind.
Let me write a complete implementation:
```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>
#include <cmath>
__device__ __forceinline__ float gelu(float x) {
constexpr float SQRT_2_PI = 0.7978845608f;
float x3 = x * x * x;
float inner = SQRT_2_PI * (x + 0.044715f * x3);
return 0.5f * x * (1.0f + tanhf(inner));
}
__global__ void compute_stats_kernel(
const float* __restrict__ input,
float* __restrict__ mean,
float* __restrict__ rstd,
int batch_size,
int channels,
int seq_len
) {
int batch_idx = blockIdx.x;
if (batch_idx >= batch_size) return;
int tid = threadIdx.x;
int numel = channels * seq_len;
extern __shared__ float sdata[];
float* s_sum = sdata;
float* s_sum_sq = sdata + blockDim.x;
float local_sum = 0.0f;
float local_sum_sq = 0.0f;
for (int i = tid; i < numel; i += blockDim.x) {
int idx = batch_idx * numel + i;
float val = input[idx];
local_sum += val;
local_sum_sq += val * val;
}
s_sum[tid] = local_sum;
s_sum_sq[tid] = local_sum_sq;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) {
s_sum[tid] += s_sum[tid + s];
s_sum_sq[tid] += s_sum_sq[tid + s];
}
__syncthreads();
}
if (tid == 0) {
float mean_val = s_sum[0] / numel;
float var_val = s_sum_sq[0] / numel - mean_val * mean_val;
mean[batch_idx] = mean_val;
rstd[batch_idx] = 1.0f / sqrtf(var_val + 1e-5f);
}
}
__global__ void fused_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
const float* __restrict__ mean,
const float* __restrict__ rstd,
float* __restrict__ output,
int batch_size,
int channels,
int seq_len,
int kernel_size,
int dilation,
int padding
) {
int batch_idx = blockIdx.z;
int channel_idx = blockIdx.y;
int seq_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (batch_idx >= batch_size || channel_idx >= channels || seq_idx >= seq_len) return;
float mean_val = mean[batch_idx];
float rstd_val = rstd[batch_idx];
int base_idx = batch_idx * channels * seq_len + channel_idx * seq_len;
int w_base = channel_idx * kernel_size;
float conv_sum = 0.0f;
for (int k = 0; k < kernel_size; k++) {
int input_seq = seq_idx - padding + k * dilation;
if (input_seq >= 0 && input_seq < seq_len) {
int idx = base_idx + input_seq;
float val = input[idx];
float norm_val = (val - mean_val) * rstd_val;
float gelu_val = gelu(norm_val);
conv_sum += gelu_val * weight[w_base + k];
}
}
float out_val = conv_sum + bias[channel_idx];
out_val += input[base_idx + seq_idx]; // residual
output[base_idx + seq_idx] = out_val;
}
std::tuple<torch::Tensor, torch::Tensor> compute_stats(torch::Tensor input) {
int batch_size = input.size(0);
int channels = input.size(1);
int seq_len = input.size(2);
auto mean = torch::empty({batch_size}, input.options());
auto rstd = torch::empty({batch_size}, input.options());
int threads = 1024;
int shared_mem = 2 * threads * sizeof(float);
compute_stats_kernel<<<batch_size, threads, shared_mem>>>(
input.data_ptr<float>(),
mean.data_ptr<float>(),
rstd.data_ptr<float>(),
batch_size, channels, seq_len
);
return std::make_tuple(mean, rstd);
}
torch::Tensor fused_norm_gelu_conv_add(
torch::Tensor input,
torch::Tensor weight,
torch::Tensor bias,
torch::Tensor mean,
torch::Tensor rstd,
int kernel_size,
int dilation,
int padding
) {
int batch_size = input.size(0);
int channels = input.size(1);
int seq_len = input.size(2);
auto output = torch::empty_like(input);
int threads = 256;
dim3 grid((seq_len + threads - 1) / threads, channels, batch_size);
fused_kernel<<<grid, threads>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
mean.data_ptr<float>(),
rstd.data_ptr<float>(),
output.data_ptr<float>(),
batch_size, channels, seq_len,
kernel_size, dilation, padding
);
return output;
}
"""
# Load the extension
fused_ops = load_inline(
name='fused_ops',
cpp_sources=cuda_source,
cuda_sources=cuda_source,
functions=['compute_stats', 'fused_norm_gelu_conv_add'],
extra_cuda_cflags=['-O3'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self, channels, seq_len, kernel_size, dilation):
super(ModelNew, self).__init__()
self.channels = channels
self.seq_len = seq_len
self.kernel_size = kernel_size
self.dilation = dilation
self.padding = dilation
# Create parameters matching nn.LayerNorm and nn.Conv1d
# LayerNorm has weight and bias for each normalized dimension
# Here we normalize over [channels, seq_len], so weight/bias shape is [channels, seq_len]
# But typical LayerNorm normalizes over last dimension(s), creating weight/bias of that shape
# Let me check: nn.LayerNorm([channels, seq_len]) normalizes over channels*seq_len
# and has weight/bias of shape [channels, seq_len]
self.norm_weight = nn.Parameter(torch.ones(channels, seq_len))
self.norm_bias = nn.Parameter(torch.zeros(channels, seq_len))
# Conv1d weight shape: [out_channels, in_channels, kernel_size]
# For groups=1 (default), this is [channels, channels, kernel_size]
# But wait, the original model has in_channels=channels, out_channels=channels
self.conv_weight = nn.Parameter(torch.randn(channels, channels, kernel_size))
self.conv_bias = nn.Parameter(torch.zeros(channels))
# Initialize conv weights similar to nn.Conv1d
nn.init.kaiming_uniform_(self.conv_weight, a=math.sqrt(5))
if self.conv_bias is not None:
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.conv_weight)
bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
nn.init.uniform_(self.conv_bias, -bound, bound)
def forward(self, x):
# Compute LayerNorm stats
mean, rstd = fused_ops.compute_stats(x)
# Fused: norm, gelu, conv, add residual
# Wait, the original Conv1d has shape [channels, channels, kernel_size]
# That means it's NOT a depthwise conv, it's a full conv!
# So each output channel depends on ALL input channels
# I need to reconsider...
pass
```
Ah, I see the issue now. The Conv1d is not depthwise - it has shape [channels, channels, kernel_size], meaning it's a full 1D convolution where each output channel is a weighted sum of all input channels. This makes the fusion much more complex because each output position needs to access all channels.
Let me reconsider the approach. The Conv1d is the main bottleneck, and it's a dense convolution. Fusing LayerNorm and GELU before the conv would require either:
1. Storing intermediate results to global memory
2. Computing on-the-fly, but with redundant computations
Actually, I think the best approach for this problem is:
1. Keep the Conv1d as-is (PyTorch's implementation is highly optimized)
2. Fuse LayerNorm + GELU into one kernel
3. Fuse residual addition at the end
But wait, if I'm going to have separate kernels anyway, let me think about what's fastest:
Option A:
- Kernel 1: LayerNorm (PyTorch)
- Kernel 2: GELU (PyTorch)
- Kernel 3: Conv1d (PyTorch, highly optimized with cuDNN)
- Kernel 4: Add residual (PyTorch)
Option B:
- Kernel 1: Fused LayerNorm + GELU
- Kernel 2: Conv1d (PyTorch, but with pre-normalized input)
- Kernel 3: Add residual
Option C:
- Kernel 1: LayerNorm (PyTorch)
- Kernel 2: GELU + Conv1d (fused, but need to compute GELU on-the-fly)
- Kernel 3: Add residual
Actually, let me think about this differently. The problem says we need to beat PyTorch's reference. Let me check what operations are happening:
1. LayerNorm: This is memory-bound (reading all data, computing stats, writing normalized data)
2. GELU: Memory-bound element-wise op
3. Conv1d: This is the expensive one - O(batch * channels * channels * seq_len * kernel_size)
4. Add residual: Memory-bound element-wise op
For a full Conv1d with [32, 256, 4096] input and [256, 256, 3] kernel:
- 32 * 256 * 4096 * 256 * 3 ≈ 25.7 billion multiply-adds
- That's a lot of compute!
The Conv1d is definitely the bottleneck, and PyTorch's implementation uses cuDNN which is highly optimized. It's hard to beat that with a custom kernel.
However, we can still fuse the surrounding operations to save memory bandwidth:
1. Fuse LayerNorm + GELU into one kernel (save one read/write)
2. Fuse residual addition at the end
Let me try this simpler approach first:
Actually, wait. Looking at the LayerNorm more carefully:
- `nn.LayerNorm([channels, seq_len])` - this normalizes over the last two dimensions together
- Input shape is [batch, channels, seq_len]
- So for each batch element, we compute mean/var over channels*seq_len = 256*4096 = 1M values
- Then apply element-wise affine transform with weight/bias of shape [channels, seq_len]
This is actually a 2D normalization, not 1D. The stats are computed per-batch element over all channels and sequence positions.
Let me write a simpler fused kernel that just handles LayerNorm + GELU + residual addition, and uses PyTorch for Conv1d:
Hmm, but the residual is added after Conv1d, not before. And Conv1d requires the GELU output as input. So we can't easily fuse everything.
Let me think about this more carefully:
```
x_norm = LayerNorm(x) # [B, C, L]
x_gelu = GELU(x_norm) # [B, C, L]
x_conv = Conv1d(x_gelu) # [B, C, L]
output = x_conv + x # [B, C, L]
```
The residual `x` is the original input, not the normalized one.
So a fused approach would be:
1. Compute LayerNorm stats for each batch element
2. For each output element, we need:
- Original input value (for residual)
- Normalized+GELU values for all channels at specific sequence positions (for Conv1d)
For Conv1d with kernel [C, C, K], each output element (b, out_c, s) needs:
- GELU values at positions (b, all_in_c, s+k*dilation) for k in [0, K)
This is complex to fuse because each output needs values from all channels.
Let me try a different approach - fuse what we can:
Fused kernel approach:
1. Compute LayerNorm stats per batch
2. For each output element:
- Read original input for residual
- Compute normalized+GELU for this position
- Compute Conv1d output by reading normalized+GELU from other positions
- Add residual
But computing Conv1d requires accessing GELU values from other positions in all channels. This is a lot of data access.
Let me try the following optimized approach:
Phase 1: Compute LayerNorm stats + transform to normalized+GELU (write to temporary buffer)
Phase 2: Conv1d (use PyTorch, but the temp buffer is already GELU output)
Phase 3: Add residual
But phase 3 needs the original input, which we still have. And phase 1 can fuse LayerNorm + GELU into one pass.
Actually, let me try fusing everything with shared memory:
The idea is:
1. Each thread block handles one (batch, out_channel, seq_group) region
2. Load stats into registers
3. Load input values into shared memory (for residual and for computing normalized+GELU on-the-fly)
4. Compute Conv1d by reading normalized+GELU from shared memory
5. Write output
The challenge is that Conv1d needs values from all input channels, so shared memory would need to hold all channels for the relevant seq positions.
With channels=256 and kernel_size=3 with dilation=4, each output position needs 3 seq positions per channel, totaling 256*3 = 768 values per output seq position per batch.
If we process multiple seq positions per block, we can share the loaded data.
Let me think about a tiling strategy:
- Each block handles (batch, out_channel, seq_tile) where seq_tile has multiple consecutive seq positions
- We load all channels for the seq positions needed by this tile
- The seq_tile size determines how much we need to extend the loaded region for the conv
For dilation=4, kernel_size=3, padding=4:
- If we process seq positions [s, s+tile_size), we need to access seq positions [s-4, s+tile_size+4] for the conv
- So we need to load seq range [s-4, s+tile_size+4], which is tile_size+8 positions
With tile_size=32 and channels=256:
- Need to load 40 seq positions * 256 channels = 10,240 floats = 40KB per block
- This fits in shared memory (228KB per SM)
But wait, for Conv1d, we also need the conv weights! The weights are [out_channels, in_channels, kernel_size].
For one out_channel, we need weights of shape [in_channels, kernel_size] = [256, 3] = 768 floats = 3KB.
So the total shared memory per block would be:
- Input values: 40 * 256 * 4 = 40KB
- Conv weights for one out_channel: 256 * 3 * 4 = 3KB
- Total: ~43KB per block
This fits! But the challenge is that loading all input channels for many seq positions takes time, and we need to do it efficiently.
Actually, let me simplify further. Let me just fuse LayerNorm + GELU into one kernel, and keep Conv1d + residual as separate. This should give a decent speedup:
```python
# Fused LayerNorm + GELU
x_norm_gelu = fused_layer_norm_gelu(x) # One kernel
# Conv1d (PyTorch optimized)
x_conv = F.conv1d(x_norm_gelu, weight, bias, padding=dilation, dilation=dilation)
# Add residual
output = x_conv + x
```
This saves one memory read/write pass compared to separate LayerNorm and GELU.
Let me implement this first, then see if I can do better.
Actually wait - I realize I need to handle the Conv1d weight shape correctly. Let me re-read the reference:
```python
self.conv = nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation)
```
So Conv1d(in_channels=channels, out_channels=channels, kernel_size=kernel_size).
With groups=1 (default), the weight shape is [out_channels, in_channels, kernel_size] = [channels, channels, kernel_size].
For each output channel oc, the output at sequence position s is:
out[oc, s] = sum over ic in [0, channels), k in [0, kernel_size):
input[ic, s - padding + k * dilation] * weight[oc, ic, k]
+ bias[oc]
This is a full convolution where each output channel depends on ALL input channels. This makes the problem more complex for fusion.
OK let me just write a clean implementation that fuses what I can, and use PyTorch for the Conv1d since it's highly optimized with cuDNN:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\nimport math\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// GELU activation function\n__device__ __forceinline__ float gelu(float x) {\n // GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))\n constexpr float SQRT_2_PI = 0.7978845608f;\n float x3 = x * x * x;\n float inner = SQRT_2_PI * (x + 0.044715f * x3);\n return 0.5f * x * (1.0f + tanhf(inner));\n}\n\n// Kernel to compute per-batch mean and sum of squares for LayerNorm\n__global__ void compute_layernorm_stats_kernel(\n const float* __restrict__ input,\n float* __restrict__ mean,\n float* __restrict__ rstd,\n int batch_size,\n int channels,\n int seq_len\n) {\n int batch_idx = blockIdx.x;\n if (batch_idx >= batch_size) return;\n \n int tid = threadIdx.x;\n int numel = channels * seq_len;\n \n extern __shared__ float sdata[];\n float* s_sum = sdata;\n float* s_sum_sq = sdata + blockDim.x;\n \n float local_sum = 0.0f;\n float local_sum_sq = 0.0f;\n \n // Strided loop over all elements in this batch\n for (int i = tid; i < numel; i += blockDim.x) {\n int idx = batch_idx * numel + i;\n float val = input[idx];\n local_sum += val;\n local_sum_sq += val * val;\n }\n \n s_sum[tid] = local_sum;\n s_sum_sq[tid] = local_sum_sq;\n __syncthreads();\n \n // Block reduction\n for (int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (tid < s) {\n s_sum[tid] += s_sum[tid + s];\n s_sum_sq[tid] += s_sum_sq[tid + s];\n }\n __syncthreads();\n }\n \n if (tid == 0) {\n float mean_val = s_sum[0] / numel;\n float var_val = s_sum_sq[0] / numel - mean_val * mean_val;\n mean[batch_idx] = mean_val;\n rstd[batch_idx] = 1.0f / sqrtf(var_val + 1e-5f);\n }\n}\n\n// Fused kernel: LayerNorm + GELU\n__global__ void layernorm_gelu_kernel(\n const float* __restrict__ input,\n const float* __restrict__ norm_weight,\n const float* __restrict__ norm_bias,\n const float* __restrict__ mean,\n const float* __restrict__ rstd,\n float* __restrict__ output,\n int batch_size,\n int channels,\n int seq_len\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total_elements = batch_size * channels * seq_len;\n \n if (idx >= total_elements) return;\n \n int batch_idx = idx / (channels * seq_len);\n int remaining = idx % (channels * seq_len);\n int channel_idx = remaining / seq_len;\n int seq_idx = remaining % seq_len;\n \n float val = input[idx];\n float mean_val = mean[batch_idx];\n float rstd_val = rstd[batch_idx];\n \n // Normalize\n float norm_val = (val - mean_val) * rstd_val;\n \n // Apply affine transform (weight and bias are [channels, seq_len])\n int norm_idx = channel_idx * seq_len + seq_idx;\n norm_val = norm_val * norm_weight[norm_idx] + norm_bias[norm_idx];\n \n // Apply GELU\n float gelu_val = gelu(norm_val);\n \n output[idx] = gelu_val;\n}\n\n// Fused kernel: Conv1d + residual add\n// Each thread handles one output element\n__global__ void conv1d_residual_kernel(\n const float* __restrict__ input_gelu,\n const float* __restrict__ residual,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n float* __restrict__ output,\n int batch_size,\n int channels,\n int seq_len,\n int kernel_size,\n int dilation,\n int padding\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total_elements = batch_size * channels * seq_len;\n \n if (idx >= total_elements) return;\n \n int batch_idx = idx / (channels * seq_len);\n int remaining = idx % (channels * seq_len);\n int out_channel = remaining / seq_len;\n int out_seq = remaining % seq_len;\n \n float conv_sum = 0.0f;\n \n // Conv1d: for each output channel, sum over all input channels\n for (int in_channel = 0; in_channel < channels; in_channel++) {\n for (int k = 0; k < kernel_size; k++) {\n int in_seq = out_seq - padding + k * dilation;\n if (in_seq >= 0 && in_seq < seq_len) {\n int input_idx = batch_idx * channels * seq_len + in_channel * seq_len + in_seq;\n int weight_idx = out_channel * channels * kernel_size + in_channel * kernel_size + k;\n conv_sum += input_gelu[input_idx] * weight[weight_idx];\n }\n }\n }\n \n // Add bias\n conv_sum += bias[out_channel];\n \n // Add residual\n int residual_idx = batch_idx * channels * seq_len + out_channel * seq_len + out_seq;\n conv_sum += residual[residual_idx];\n \n output[idx] = conv_sum;\n}\n\nstd::tuple<torch::Tensor, torch::Tensor> compute_layernorm_stats(torch::Tensor input) {\n int batch_size = input.size(0);\n int channels = input.size(1);\n int seq_len = input.size(2);\n \n auto options = input.options();\n auto mean = torch::empty({batch_size}, options);\n auto rstd = torch::empty({batch_size}, options);\n \n int threads = 1024;\n int shared_mem = 2 * threads * sizeof(float);\n \n compute_layernorm_stats_kernel<<<batch_size, threads, shared_mem>>>(\n input.data_ptr<float>(),\n mean.data_ptr<float>(),\n rstd.data_ptr<float>(),\n batch_size, channels, seq_len\n );\n \n return std::make_tuple(mean, rstd);\n}\n\ntorch::Tensor layernorm_gelu(\n torch::Tensor input,\n torch::Tensor norm_weight,\n torch::Tensor norm_bias,\n torch::Tensor mean,\n torch::Tensor rstd\n) {\n int batch_size = input.size(0);\n int channels = input.size(1);\n int seq_len = input.size(2);\n \n auto output = torch::empty_like(input);\n \n int total_elements = batch_size * channels * seq_len;\n int threads = 256;\n int blocks = (total_elements + threads - 1) / threads;\n \n layernorm_gelu_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n norm_weight.data_ptr<float>(),\n norm_bias.data_ptr<float>(),\n mean.data_ptr<float>(),\n rstd.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size, channels, seq_len\n );\n \n return output;\n}\n\ntorch::Tensor conv1d_residual(\n torch::Tensor input_gelu,\n torch::Tensor residual,\n torch::Tensor weight,\n torch::Tensor bias,\n int kernel_size,\n int dilation,\n int padding\n) {\n int batch_size = input_gelu.size(0);\n int channels = input_gelu.size(1);\n int seq_len = input_gelu.size(2);\n \n auto output = torch::empty_like(input_gelu);\n \n int total_elements = batch_size * channels * seq_len;\n int threads = 256;\n int blocks = (total_elements + threads - 1) / threads;\n \n conv1d_residual_kernel<<<blocks, threads>>>(\n input_gelu.data_ptr<float>(),\n residual.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size, channels, seq_len,\n kernel_size, dilation, padding\n );\n \n return output;\n}\n\"\"\"\n\n# Load the extension\nfused_ops = load_inline(\n name='fused_ops_v1',\n cpp_sources=cuda_source,\n cuda_sources=cuda_source,\n functions=['compute_layernorm_stats', 'layernorm_gelu', 'conv1d_residual'],\n extra_cuda_cflags=['-O3'],\n verbose=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, channels, seq_len, kernel_size, dilation):\n super(ModelNew, self).__init__()\n self.channels = channels\n self.seq_len = seq_len\n self.kernel_size = kernel_size\n self.dilation = dilation\n self.padding = dilation\n \n # LayerNorm parameters: normalized_shape=[channels, seq_len]\n # So weight and bias have shape [channels, seq_len]\n self.norm_weight = nn.Parameter(torch.ones(channels, seq_len))\n self.norm_bias = nn.Parameter(torch.zeros(channels, seq_len))\n \n # Conv1d parameters: [out_channels, in_channels, kernel_size]\n # With in_channels=channels, out_channels=channels\n self.conv_weight = nn.Parameter(torch.empty(channels, channels, kernel_size))\n self.conv_bias = nn.Parameter(torch.zeros(channels))\n \n # Initialize conv weights similar to PyTorch\n nn.init.kaiming_uniform_(self.conv_weight, a=math.sqrt(5))\n fan_in = channels * kernel_size\n bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0\n nn.init.uniform_(self.conv_bias, -bound, bound)\n \n def forward(self, x):\n # Compute LayerNorm statistics\n mean, rstd = fused_ops.compute_layernorm_stats(x)\n \n # Fused LayerNorm + GELU\n x_gelu = fused_ops.layernorm_gelu(x, self.norm_weight, self.norm_bias, mean, rstd)\n \n # Fused Conv1d + residual\n output = fused_ops.conv1d_residual(\n x_gelu, x, self.conv_weight, self.conv_bias,\n self.kernel_size, self.dilation, self.padding\n )\n \n return output"
}