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# Fused operations for Hyena\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <c10/cuda/CUDAGuard.h>\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n\n// Fused projection split and transpose: (B, L, D*(order+1)) -> multiple (B, D, L) outputs\n// This avoids the intermediate contiguous/chunk operations\n__global__ void fused_proj_split_transpose_kernel(\n const float* __restrict__ projected,\n float* __restrict__ v_out,\n float* __restrict__* gates_out,\n int batch_size,\n int seq_len,\n int d_model,\n int order\n) {\n __shared__ float tile[32][33];\n \n int b = blockIdx.z;\n int tile_seq = blockIdx.y * 32;\n int tile_d = blockIdx.x * 32;\n \n int seq_local = threadIdx.y;\n int d_local = threadIdx.x;\n \n int seq_global = tile_seq + seq_local;\n int d_global = tile_d + d_local;\n \n int total_d = d_model * (order + 1);\n \n // Load v component (first d_model dimensions)\n if (seq_global < seq_len && d_global < d_model && b < batch_size) {\n tile[d_local][seq_local] = projected[b * seq_len * total_d + seq_global * total_d + d_global];\n }\n __syncthreads();\n \n int out_seq = tile_seq + d_local;\n int out_d = tile_d + seq_local;\n \n if (out_seq < seq_len && out_d < d_model && b < batch_size) {\n v_out[b * d_model * seq_len + out_d * seq_len + out_seq] = tile[seq_local][d_local];\n }\n \n // Now handle each gate\n for (int g = 0; g < order; g++) {\n __syncthreads();\n int gate_d_offset = d_model * (g + 1);\n \n if (seq_global < seq_len && d_global < d_model && b < batch_size) {\n tile[d_local][seq_local] = projected[b * seq_len * total_d + seq_global * total_d + gate_d_offset + d_global];\n }\n __syncthreads();\n \n if (out_seq < seq_len && out_d < d_model && b < batch_size) {\n gates_out[g][b * d_model * seq_len + out_d * seq_len + out_seq] = tile[seq_local][d_local];\n }\n }\n}\n\n// Optimized depthwise conv1d\n__global__ void depthwise_conv1d_optimized(\n const float* __restrict__ input,\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) {\n extern __shared__ float shared_mem[];\n float* shared_input = shared_mem;\n \n int b = blockIdx.y;\n int c = blockIdx.z;\n int s_start = blockIdx.x * blockDim.x;\n int s_local = threadIdx.x;\n int s = s_start + s_local;\n int half_k = kernel_size / 2;\n \n float sum = (bias != nullptr) ? bias[c] : 0.0f;\n \n int pad = half_k;\n for (int i = s_local; i < blockDim.x + 2 * pad; i += blockDim.x) {\n int s_global = s_start + i - pad;\n if (s_global >= 0 && s_global < seq_len) {\n shared_input[i] = input[b * channels * seq_len + c * seq_len + s_global];\n } else {\n shared_input[i] = 0.0f;\n }\n }\n __syncthreads();\n \n if (s < seq_len && b < batch_size && c < channels) {\n for (int k = 0; k < kernel_size; k++) {\n sum += shared_input[s_local + pad + k - half_k] * weight[c * kernel_size + k];\n }\n output[b * channels * seq_len + c * seq_len + s] = sum;\n }\n}\n\n// Fused gating and transpose: applies gating then transposes (B, D, L) -> (B, L, D)\n__global__ void fused_gated_transpose_kernel(\n float* __restrict__ y,\n const float* __restrict__ gate,\n float* __restrict__ output,\n int batch_size,\n int d_model,\n int seq_len\n) {\n __shared__ float tile[32][33];\n \n int b = blockIdx.z;\n int tile_d = blockIdx.y * 32;\n int tile_seq = blockIdx.x * 32;\n \n int d_local = threadIdx.y;\n int seq_local = threadIdx.x;\n \n int d_global = tile_d + d_local;\n int seq_global = tile_seq + seq_local;\n \n // Load, apply gate, and store to shared memory\n if (d_global < d_model && seq_global < seq_len && b < batch_size) {\n int idx = b * d_model * seq_len + d_global * seq_len + seq_global;\n tile[seq_local][d_local] = y[idx] * gate[idx];\n y[idx] = 0.0f; // Clear for potential reuse\n }\n __syncthreads();\n \n // Write transposed to output\n int out_d = tile_d + seq_local;\n int out_seq = tile_seq + d_local;\n \n if (out_d < d_model && out_seq < seq_len && b < batch_size) {\n output[b * seq_len * d_model + out_seq * d_model + out_d] = tile[d_local][seq_local];\n }\n}\n\n// Simple transpose for final output\n__global__ void transpose_bdl_bld_kernel(\n const float* __restrict__ input,\n float* __restrict__ output,\n int batch_size,\n int d_model,\n int seq_len\n) {\n __shared__ float tile[32][33];\n \n int b = blockIdx.z;\n int tile_d = blockIdx.y * 32;\n int tile_seq = blockIdx.x * 32;\n \n int d_local = threadIdx.y;\n int seq_local = threadIdx.x;\n \n int d_global = tile_d + d_local;\n int seq_global = tile_seq + seq_local;\n \n if (d_global < d_model && seq_global < seq_len && b < batch_size) {\n tile[seq_local][d_local] = input[b * d_model * seq_len + d_global * seq_len + seq_global];\n }\n __syncthreads();\n \n int out_d = tile_d + seq_local;\n int out_seq = tile_seq + d_local;\n \n if (out_d < d_model && out_seq < seq_len && b < batch_size) {\n output[b * seq_len * d_model + out_seq * d_model + out_d] = tile[d_local][seq_local];\n }\n}\n\n// Standalone elementwise multiplication\n__global__ void elementwise_mul_kernel(\n float* __restrict__ data,\n const float* __restrict__ gate,\n int total_elements\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total_elements) {\n data[idx] *= gate[idx];\n }\n}\n\n// Python interface functions\nstd::vector<torch::Tensor> fused_proj_split_transpose(\n torch::Tensor projected,\n int d_model,\n int order\n) {\n CHECK_CUDA(projected);\n \n int batch_size = projected.size(0);\n int seq_len = projected.size(1);\n \n auto v_out = torch::empty({batch_size, d_model, seq_len}, projected.options());\n std::vector<torch::Tensor> gates;\n std::vector<float*> gate_ptrs;\n for (int i = 0; i < order; i++) {\n gates.push_back(torch::empty({batch_size, d_model, seq_len}, projected.options()));\n gate_ptrs.push_back(gates[i].data_ptr<float>());\n }\n \n dim3 block(32, 32);\n dim3 grid((d_model + 31) / 32, (seq_len + 31) / 32, batch_size);\n \n // Need to copy gate_ptrs to device accessible memory\n thrust::device_vector<float*> d_gate_ptrs(gate_ptrs);\n \n fused_proj_split_transpose_kernel<<<grid, block>>>(\n projected.data_ptr<float>(),\n v_out.data_ptr<float>(),\n thrust::raw_pointer_cast(d_gate_ptrs.data()),\n batch_size, seq_len, d_model, order\n );\n \n gates.insert(gates.begin(), v_out);\n return gates;\n}\n\ntorch::Tensor depthwise_conv1d(\n torch::Tensor input,\n torch::Tensor weight,\n torch::Tensor bias,\n int kernel_size\n) {\n CHECK_CUDA(input);\n CHECK_CUDA(weight);\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 threads = 256;\n dim3 grid((seq_len + threads - 1) / threads, batch_size, channels);\n int shared_mem = (threads + kernel_size - 1) * sizeof(float);\n \n const float* bias_ptr = bias.defined() && bias.numel() > 0 ? bias.data_ptr<float>() : nullptr;\n \n depthwise_conv1d_optimized<<<grid, threads, shared_mem>>>(\n input.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias_ptr,\n output.data_ptr<float>(),\n batch_size, channels, seq_len, kernel_size\n );\n \n return output;\n}\n\ntorch::Tensor fused_gated_transpose(\n torch::Tensor y,\n torch::Tensor gate\n) {\n CHECK_CUDA(y);\n CHECK_CUDA(gate);\n \n int batch_size = y.size(0);\n int d_model = y.size(1);\n int seq_len = y.size(2);\n \n auto output = torch::empty({batch_size, seq_len, d_model}, y.options());\n \n dim3 block(32, 32);\n dim3 grid((seq_len + 31) / 32, (d_model + 31) / 32, batch_size);\n \n fused_gated_transpose_kernel<<<grid, block>>>(\n y.data_ptr<float>(),\n gate.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size, d_model, seq_len\n );\n \n return output;\n}\n\ntorch::Tensor transpose_bdl_bld(torch::Tensor input) {\n CHECK_CUDA(input);\n \n int batch_size = input.size(0);\n int d_model = input.size(1);\n int seq_len = input.size(2);\n \n auto output = torch::empty({batch_size, seq_len, d_model}, input.options());\n \n dim3 block(32, 32);\n dim3 grid((seq_len + 31) / 32, (d_model + 31) / 32, batch_size);\n \n transpose_bdl_bld_kernel<<<grid, block>>>(\n input.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size, d_model, seq_len\n );\n \n return output;\n}\n\ntorch::Tensor elementwise_mul(torch::Tensor data, torch::Tensor gate) {\n CHECK_CUDA(data);\n CHECK_CUDA(gate);\n \n int total_elements = data.numel();\n \n int threads = 256;\n int blocks = (total_elements + threads - 1) / threads;\n \n elementwise_mul_kernel<<<blocks, threads>>>(\n data.data_ptr<float>(),\n gate.data_ptr<float>(),\n total_elements\n );\n \n return data;\n}\n\"\"\"\n\ncpp_source = \"\"\"\n#include <thrust/device_vector.h>\n\nstd::vector<torch::Tensor> fused_proj_split_transpose(torch::Tensor projected, int d_model, int order);\ntorch::Tensor depthwise_conv1d(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, int kernel_size);\ntorch::Tensor fused_gated_transpose(torch::Tensor y, torch::Tensor gate);\ntorch::Tensor transpose_bdl_bld(torch::Tensor input);\ntorch::Tensor elementwise_mul(torch::Tensor data, torch::Tensor gate);\n\"\"\"\n\nhyena_extension = load_inline(\n name='hyena_extension',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['fused_proj_split_transpose', 'depthwise_conv1d', 'fused_gated_transpose', 'transpose_bdl_bld', 'elementwise_mul'],\n extra_cuda_cflags=['-O3', '-use_fast_math', '--expt-relaxed-constexpr'],\n extra_ldflags=['-lcudart'],\n verbose=True\n)\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, seq_len: int, order: int, short_filter_size: int):\n super(ModelNew, self).__init__()\n self.d_model = d_model\n self.seq_len = seq_len\n self.order = order\n self.fft_size = 2 * seq_len\n self.short_filter_size = short_filter_size\n\n self.in_proj = nn.Linear(d_model, d_model * (order + 1), bias=True)\n\n self.short_convs = nn.ModuleList([\n nn.Conv1d(d_model, d_model, short_filter_size,\n padding=short_filter_size // 2, groups=d_model)\n for _ in range(order)\n ])\n\n self.filter_mlp = nn.Sequential(\n nn.Linear(1, 64),\n nn.SiLU(),\n nn.Linear(64, d_model),\n )\n t = torch.linspace(0, 1, seq_len).unsqueeze(-1)\n self.register_buffer('filter_positions', t)\n\n self.out_proj = nn.Linear(d_model, d_model, bias=True)\n \n # Pre-compute filter\n with torch.no_grad():\n h = self.filter_mlp(self.filter_positions).transpose(0, 1)\n h_padded = F.pad(h, (0, seq_len))\n self.register_buffer('h_f', torch.fft.rfft(h_padded, dim=-1))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, L, D = x.shape\n \n # Linear projection\n projected = self.in_proj(x)\n \n # Fused split and transpose\n split_tensors = hyena_extension.fused_proj_split_transpose(projected.contiguous(), self.d_model, self.order)\n v = split_tensors[0]\n gates = split_tensors[1:]\n \n # Depthwise convolutions\n for i in range(self.order):\n conv = self.short_convs[i]\n gates[i] = hyena_extension.depthwise_conv1d(\n gates[i],\n conv.weight.squeeze(2).contiguous(),\n conv.bias if conv.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype),\n self.short_filter_size\n )\n\n # FFT convolution\n v_padded = F.pad(v, (0, self.seq_len))\n v_f = torch.fft.rfft(v_padded, dim=-1)\n y_f = v_f * self.h_f.unsqueeze(0)\n y = torch.fft.irfft(y_f, n=self.fft_size, dim=-1)[..., :L]\n\n # Gating\n for gate in gates:\n y = hyena_extension.elementwise_mul(y, gate)\n \n # Final transpose\n y = hyena_extension.transpose_bdl_bld(y)\n y = self.out_proj(y)\n return y"
}