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\nimport math\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cufft.h>\n#include <cuda_fp16.h>\n#include <cublas_v2.h>\n\n#define CUFFT_CALL(call) \\\n do { \\\n cufftResult_t err = call; \\\n if (err != CUFFT_SUCCESS) { \\\n printf(\"cuFFT error %d at %s:%d\\n\", err, __FILE__, __LINE__); \\\n } \\\n } while (0)\n\n// Kernel for complex multiplication: y = v * h (batched)\n__global__ void complex_multiply_kernel(\n const cufftComplex* __restrict__ v_f,\n const cufftComplex* __restrict__ h_f,\n cufftComplex* __restrict__ y_f,\n int batch_size,\n int d_model,\n int fft_size\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = batch_size * d_model * fft_size;\n \n if (idx < total) {\n int b = idx / (d_model * fft_size);\n int rem = idx % (d_model * fft_size);\n int d = rem / fft_size;\n int f = rem % fft_size;\n \n int h_idx = d * fft_size + f;\n int v_idx = idx;\n \n cufftComplex v_val = v_f[v_idx];\n cufftComplex h_val = h_f[h_idx];\n \n // Complex multiplication\n cufftComplex result;\n result.x = v_val.x * h_val.x - v_val.y * h_val.y;\n result.y = v_val.x * h_val.y + v_val.y * h_val.x;\n \n y_f[idx] = result;\n }\n}\n\n// Kernel to apply gating: y = y * gate\n__global__ void apply_gate_kernel(\n float* __restrict__ y,\n const float* __restrict__ gate,\n int batch_size,\n int d_model,\n int seq_len\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = batch_size * d_model * seq_len;\n \n if (idx < total) {\n y[idx] = y[idx] * gate[idx];\n }\n}\n\n// Copy and scale kernel for IRFFT output\n__global__ void copy_scale_kernel(\n const float* __restrict__ src,\n float* __restrict__ dst,\n int batch_size,\n int d_model,\n int seq_len,\n int fft_size,\n float scale\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = batch_size * d_model * seq_len;\n \n if (idx < total) {\n int b = idx / (d_model * seq_len);\n int rem = idx % (d_model * seq_len);\n int d = rem / seq_len;\n int s = rem % seq_len;\n \n int src_idx = b * d_model * fft_size + d * fft_size + s;\n dst[idx] = src[src_idx] * scale;\n }\n}\n\n// Global cuFFT handles\nstatic cufftHandle rfft_plan = 0;\nstatic cufftHandle irfft_plan = 0;\nstatic int cached_fft_size = 0;\nstatic bool plans_initialized = false;\n\nvoid ensure_fft_plans(int fft_size) {\n if (plans_initialized && cached_fft_size == fft_size) {\n return;\n }\n \n if (plans_initialized) {\n CUFFT_CALL(cufftDestroy(rfft_plan));\n CUFFT_CALL(cufftDestroy(irfft_plan));\n }\n \n // R2C plan: input size fft_size, batch size d_model * batch_size\n // We'll create plans for batched transforms\n CUFFT_CALL(cufftPlan1d(&rfft_plan, fft_size, CUFFT_R2C, 1));\n CUFFT_CALL(cufftPlan1d(&irfft_plan, fft_size, CUFFT_C2R, 1));\n \n cached_fft_size = fft_size;\n plans_initialized = true;\n}\n\n// Main fused kernel for FFT convolution + gating\nvoid hyena_fft_conv(\n const float* v_ptr,\n const float* h_ptr,\n const float* gate0_ptr,\n const float* gate1_ptr,\n float* y_ptr,\n int batch_size,\n int d_model,\n int seq_len,\n int fft_size,\n cudaStream_t stream\n) {\n ensure_fft_plans(fft_size);\n \n int rfft_size = fft_size / 2 + 1;\n int total_elements = batch_size * d_model * fft_size;\n int total_complex = batch_size * d_model * rfft_size;\n \n // Allocate workspace\n float* v_padded;\n float* h_padded;\n cufftComplex* v_f;\n cufftComplex* h_f;\n cufftComplex* y_f;\n float* y_padded;\n \n cudaMalloc(&v_padded, total_elements * sizeof(float));\n cudaMalloc(&h_padded, d_model * fft_size * sizeof(float));\n cudaMalloc(&v_f, total_complex * sizeof(cufftComplex));\n cudaMalloc(&h_f, d_model * rfft_size * sizeof(cufftComplex));\n cudaMalloc(&y_f, total_complex * sizeof(cufftComplex));\n cudaMalloc(&y_padded, total_elements * sizeof(float));\n \n // Zero out padded arrays\n cudaMemset(v_padded, 0, total_elements * sizeof(float));\n cudaMemset(h_padded, 0, d_model * fft_size * sizeof(float));\n \n // Copy v to padded array (v is B x D x L, we need B x D x fft_size)\n // Copy each batch/channel\n for (int b = 0; b < batch_size; b++) {\n for (int d = 0; d < d_model; d++) {\n int src_offset = b * d_model * seq_len + d * seq_len;\n int dst_offset = b * d_model * fft_size + d * fft_size;\n cudaMemcpyAsync(\n v_padded + dst_offset,\n v_ptr + src_offset,\n seq_len * sizeof(float),\n cudaMemcpyDeviceToDevice,\n stream\n );\n }\n }\n \n // Copy h to padded array (h is D x L)\n for (int d = 0; d < d_model; d++) {\n int src_offset = d * seq_len;\n int dst_offset = d * fft_size;\n cudaMemcpyAsync(\n h_padded + dst_offset,\n h_ptr + src_offset,\n seq_len * sizeof(float),\n cudaMemcpyDeviceToDevice,\n stream\n );\n }\n \n // Execute R2C FFT for v (batched over B*D)\n cufftSetStream(rfft_plan, stream);\n for (int b = 0; b < batch_size; b++) {\n for (int d = 0; d < d_model; d++) {\n int input_offset = b * d_model * fft_size + d * fft_size;\n int output_offset = b * d_model * rfft_size + d * rfft_size;\n CUFFT_CALL(cufftExecR2C(\n rfft_plan,\n (cufftReal*)(v_padded + input_offset),\n v_f + output_offset\n ));\n }\n }\n \n // Execute R2C FFT for h\n for (int d = 0; d < d_model; d++) {\n int input_offset = d * fft_size;\n int output_offset = d * rfft_size;\n CUFFT_CALL(cufftExecR2C(\n rfft_plan,\n (cufftReal*)(h_padded + input_offset),\n h_f + output_offset\n ));\n }\n \n // Complex multiply: y_f = v_f * h_f\n int threads = 256;\n int blocks = (total_complex + threads - 1) / threads;\n complex_multiply_kernel<<<blocks, threads, 0, stream>>>(\n v_f, h_f, y_f, batch_size, d_model, rfft_size\n );\n \n // Execute C2R IFFT for y\n cufftSetStream(irfft_plan, stream);\n for (int b = 0; b < batch_size; b++) {\n for (int d = 0; d < d_model; d++) {\n int input_offset = b * d_model * rfft_size + d * rfft_size;\n int output_offset = b * d_model * fft_size + d * fft_size;\n CUFFT_CALL(cufftExecC2R(\n irfft_plan,\n y_f + input_offset,\n y_padded + output_offset\n ));\n }\n }\n \n // Copy first L elements and scale\n float scale = 1.0f / fft_size;\n blocks = (batch_size * d_model * seq_len + threads - 1) / threads;\n copy_scale_kernel<<<blocks, threads, 0, stream>>>(\n y_padded, y_ptr, batch_size, d_model, seq_len, fft_size, scale\n );\n \n // Apply gates\n apply_gate_kernel<<<blocks, threads, 0, stream>>>(\n y_ptr, gate0_ptr, batch_size, d_model, seq_len\n );\n apply_gate_kernel<<<blocks, threads, 0, stream>>>(\n y_ptr, gate1_ptr, batch_size, d_model, seq_len\n );\n \n // Cleanup\n cudaFree(v_padded);\n cudaFree(h_padded);\n cudaFree(v_f);\n cudaFree(h_f);\n cudaFree(y_f);\n cudaFree(y_padded);\n}\n\n// Entry point called from Python\nvoid hyena_forward_cuda(\n torch::Tensor v,\n torch::Tensor h,\n torch::Tensor gate0,\n torch::Tensor gate1,\n torch::Tensor y,\n int fft_size\n) {\n int batch_size = v.size(0);\n int d_model = v.size(1);\n int seq_len = v.size(2);\n \n cudaStream_t stream = c10::cuda::getCurrentCUDAStream();\n \n hyena_fft_conv(\n v.data_ptr<float>(),\n h.data_ptr<float>(),\n gate0.data_ptr<float>(),\n gate1.data_ptr<float>(),\n y.data_ptr<float>(),\n batch_size,\n d_model,\n seq_len,\n fft_size,\n stream\n );\n}\n'''\n\ncpp_source = r'''\nvoid hyena_forward_cuda(\n torch::Tensor v,\n torch::Tensor h,\n torch::Tensor gate0,\n torch::Tensor gate1,\n torch::Tensor y,\n int fft_size\n);\n\nvoid hyena_forward(\n torch::Tensor v,\n torch::Tensor h,\n torch::Tensor gate0,\n torch::Tensor gate1,\n torch::Tensor y,\n int fft_size\n) {\n hyena_forward_cuda(v, h, gate0, gate1, y, fft_size);\n}\n'''\n\nhyena_module = load_inline(\n name='hyena_ops',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['hyena_forward'],\n extra_cuda_cflags=['-O3', '-use_fast_math'],\n verbose=True\n)\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\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 def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, L, D = x.shape\n\n projected = self.in_proj(x)\n chunks = projected.chunk(self.order + 1, dim=-1)\n v = chunks[0].transpose(1, 2)\n gates = [c.transpose(1, 2) for c in chunks[1:]]\n\n for i in range(self.order):\n gates[i] = self.short_convs[i](gates[i])\n\n h = self.filter_mlp(self.filter_positions).transpose(0, 1)\n\n # Use custom CUDA kernel for FFT convolution + gating\n y = torch.empty(B, D, L, device=x.device, dtype=x.dtype)\n hyena_module.hyena_forward(v, h, gates[0], gates[1], y, self.fft_size)\n\n y = y.transpose(1, 2)\n y = self.out_proj(y)\n return y"
}