reasoning
Good, correctness passed! But the current implementation just uses PyTorch's built-in operations. I need to implement the actual custom CUDA kernel to get speedup. Let me think about the architecture:
1. The model does two attention passes: row-wise and column-wise
2. Each attention pass involves:
- LayerNorm
- QKV projection
- Attention (Q@K^T, softmax, @V)
- Output projection
For the custom CUDA kernel, I need to:
1. Efficiently handle the batch structure
2. Fuse operations where possible
3. Use shared memory effectively
4. Handle the column attention which requires different memory access patterns
Let me think about the key optimizations:
1. Fused LayerNorm + QKV projection
2. Fused attention computation
3. Using shared memory for Q, K, V
4. Efficient softmax computation
The challenge is that each "batch" element in the reshaped view is actually either a row or column of the original input. For row attention, we have B*H sequences of length W. For column attention, we have B*W sequences of length H.
Since W=32 and H=32 are fixed, I can use these as compile-time constants for optimization.
Let me design a proper kernel that:
1. Has each thread block handle one (batch_sequence, head) pair
2. Uses shared memory to load the input sequence, compute QKV, attention, and output
3. Handles the output projection properly by having all heads write to shared memory
Actually, let me think more carefully. The problem is that multiple thread blocks are working on the same sequence (one per head), and they need to combine their outputs for the output projection.
One approach: have one block handle ALL heads for one sequence. This means we have B_outer blocks, and each block processes all num_heads heads.
For our case:
- dim = 64
- num_heads = 4
- head_dim = 16
- W = H = 32
So for row attention, we have B*H = 4*32 = 128 sequences, each of length 32 and dim 64.
Each block needs to:
1. Load sequence of shape [32, 64]
2. Apply LayerNorm
3. Compute QKV: [32, 64] -> [32, 3*64]
4. Split into 4 heads, each [32, 16] for Q, K, V
5. For each head, compute attention scores [32, 32]
6. Apply softmax
7. Multiply by V to get [32, 16] per head
8. Concatenate heads to [32, 64]
9. Apply output projection to [32, 64]
Shared memory required:
- x: [32, 64] = 2048 floats
- x_norm: [32, 64] = 2048 floats
- qkv: [32, 192] = 6144 floats (or we can compute per-head)
- attn_scores: [4, 32, 32] = 4096 floats (for all heads)
- attn_out: [32, 64] = 2048 floats
Total: ~16KB which fits easily in shared memory.
Let me rewrite the kernel with this approach.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom torch.utils.cpp_extension import load_inline\n\n# CUDA source for fused attention operations\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <c10/cuda/CUDAException.h>\n#include <cmath>\n\nconstexpr int WARP_SIZE = 32;\n\n// Fused attention kernel: LayerNorm + QKV + Attention + Output projection\n// Each block handles one sequence, processing all heads\ntemplate <typename scalar_t, int SEQ_LEN, int DIM, int NUM_HEADS, int HEAD_DIM>\n__global__ void fused_attention_kernel(\n const scalar_t* __restrict__ x, // [batch, seq_len, dim]\n const scalar_t* __restrict__ norm_w, // [dim]\n const scalar_t* __restrict__ norm_b, // [dim]\n const scalar_t* __restrict__ qkv_w, // [3*dim, dim]\n const scalar_t* __restrict__ out_w, // [dim, dim]\n scalar_t* __restrict__ out, // [batch, seq_len, dim]\n int batch_size\n) {\n // Shared memory layout\n extern __shared__ char smem[];\n scalar_t* x_smem = reinterpret_cast<scalar_t*>(smem); // [seq_len, dim]\n scalar_t* x_norm = x_smem + SEQ_LEN * DIM; // [seq_len, dim]\n scalar_t* q_smem = x_norm + SEQ_LEN * DIM; // [seq_len, dim]\n scalar_t* k_smem = q_smem + SEQ_LEN * DIM; // [seq_len, dim]\n scalar_t* v_smem = k_smem + SEQ_LEN * DIM; // [seq_len, dim]\n scalar_t* attn_scores = v_smem + SEQ_LEN * DIM; // [NUM_HEADS, seq_len, seq_len]\n scalar_t* attn_out = attn_scores + NUM_HEADS * SEQ_LEN * SEQ_LEN; // [seq_len, dim]\n \n const int batch_idx = blockIdx.x;\n const int tid = threadIdx.x;\n const int blockDimX = blockDim.x;\n \n const scalar_t* x_batch = x + batch_idx * SEQ_LEN * DIM;\n scalar_t* out_batch = out + batch_idx * SEQ_LEN * DIM;\n \n // Load input into shared memory\n for (int i = tid; i < SEQ_LEN * DIM; i += blockDimX) {\n x_smem[i] = x_batch[i];\n }\n __syncthreads();\n \n // LayerNorm for each position\n for (int s = 0; s < SEQ_LEN; s++) {\n scalar_t sum = 0.0f;\n scalar_t sum_sq = 0.0f;\n \n for (int d = tid; d < DIM; d += blockDimX) {\n scalar_t val = x_smem[s * DIM + d];\n sum += val;\n sum_sq += val * val;\n }\n \n // Warp reduction (assuming blockDimX <= 128 and is power of 2)\n for (int offset = blockDimX / 2; offset > 0; offset >>= 1) {\n sum += __shfl_down_sync(0xffffffff, sum, offset);\n sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);\n }\n \n scalar_t mean = __shfl_sync(0xffffffff, sum, 0) / DIM;\n scalar_t var = __shfl_sync(0xffffffff, sum_sq, 0) / DIM - mean * mean;\n scalar_t inv_std = rsqrtf(var + 1e-5f);\n \n for (int d = tid; d < DIM; d += blockDimX) {\n x_norm[s * DIM + d] = (x_smem[s * DIM + d] - mean) * inv_std * norm_w[d] + norm_b[d];\n }\n }\n __syncthreads();\n \n // QKV projection\n for (int s = 0; s < SEQ_LEN; s++) {\n // Compute Q\n for (int d = tid; d < DIM; d += blockDimX) {\n scalar_t sum = 0.0f;\n for (int in_d = 0; in_d < DIM; in_d++) {\n sum += x_norm[s * DIM + in_d] * qkv_w[d * DIM + in_d];\n }\n q_smem[s * DIM + d] = sum;\n }\n \n // Compute K\n for (int d = tid; d < DIM; d += blockDimX) {\n scalar_t sum = 0.0f;\n for (int in_d = 0; in_d < DIM; in_d++) {\n sum += x_norm[s * DIM + in_d] * qkv_w[(DIM + d) * DIM + in_d];\n }\n k_smem[s * DIM + d] = sum;\n }\n \n // Compute V\n for (int d = tid; d < DIM; d += blockDimX) {\n scalar_t sum = 0.0f;\n for (int in_d = 0; in_d < DIM; in_d++) {\n sum += x_norm[s * DIM + in_d] * qkv_w[(2 * DIM + d) * DIM + in_d];\n }\n v_smem[s * DIM + d] = sum;\n }\n }\n __syncthreads();\n \n // Attention for each head\n const scalar_t scale = 1.0f / sqrtf((scalar_t)HEAD_DIM);\n \n for (int head = 0; head < NUM_HEADS; head++) {\n for (int q_pos = 0; q_pos < SEQ_LEN; q_pos++) {\n scalar_t max_score = -INFINITY;\n \n // Compute attention scores for this query position\n for (int k_pos = tid; k_pos < SEQ_LEN; k_pos += blockDimX) {\n scalar_t score = 0.0f;\n for (int d = 0; d < HEAD_DIM; d++) {\n score += q_smem[q_pos * DIM + head * HEAD_DIM + d] * \n k_smem[k_pos * DIM + head * HEAD_DIM + d];\n }\n score *= scale;\n attn_scores[head * SEQ_LEN * SEQ_LEN + q_pos * SEQ_LEN + k_pos] = score;\n if (score > max_score) max_score = score;\n }\n \n // Reduce max_score across threads\n for (int offset = blockDimX / 2; offset > 0; offset >>= 1) {\n max_score = fmaxf(max_score, __shfl_down_sync(0xffffffff, max_score, offset));\n }\n max_score = __shfl_sync(0xffffffff, max_score, 0);\n \n // Compute exp and sum\n scalar_t sum_exp = 0.0f;\n for (int k_pos = tid; k_pos < SEQ_LEN; k_pos += blockDimX) {\n scalar_t exp_val = expf(attn_scores[head * SEQ_LEN * SEQ_LEN + q_pos * SEQ_LEN + k_pos] - max_score);\n attn_scores[head * SEQ_LEN * SEQ_LEN + q_pos * SEQ_LEN + k_pos] = exp_val;\n sum_exp += exp_val;\n }\n \n // Reduce sum_exp\n for (int offset = blockDimX / 2; offset > 0; offset >>= 1) {\n sum_exp += __shfl_down_sync(0xffffffff, sum_exp, offset);\n }\n sum_exp = __shfl_sync(0xffffffff, sum_exp, 0);\n \n // Normalize softmax\n scalar_t inv_sum = 1.0f / sum_exp;\n for (int k_pos = tid; k_pos < SEQ_LEN; k_pos += blockDimX) {\n attn_scores[head * SEQ_LEN * SEQ_LEN + q_pos * SEQ_LEN + k_pos] *= inv_sum;\n }\n \n __syncthreads();\n \n // Weighted sum of values\n for (int d = tid; d < HEAD_DIM; d += blockDimX) {\n scalar_t val = 0.0f;\n for (int k_pos = 0; k_pos < SEQ_LEN; k_pos++) {\n val += attn_scores[head * SEQ_LEN * SEQ_LEN + q_pos * SEQ_LEN + k_pos] * \n v_smem[k_pos * DIM + head * HEAD_DIM + d];\n }\n attn_out[q_pos * DIM + head * HEAD_DIM + d] = val;\n }\n \n __syncthreads();\n }\n }\n \n // Output projection\n for (int s = 0; s < SEQ_LEN; s++) {\n for (int d = tid; d < DIM; d += blockDimX) {\n scalar_t sum = 0.0f;\n for (int in_d = 0; in_d < DIM; in_d++) {\n sum += attn_out[s * DIM + in_d] * out_w[d * DIM + in_d];\n }\n out_batch[s * DIM + d] = sum;\n }\n }\n}\n\n// Host function to launch the kernel\ntorch::Tensor fused_attention_cuda(\n torch::Tensor x,\n torch::Tensor norm_weight,\n torch::Tensor norm_bias,\n torch::Tensor qkv_weight,\n torch::Tensor out_weight,\n int num_heads,\n int seq_len,\n int dim\n) {\n int batch = x.size(0);\n int head_dim = dim / num_heads;\n \n auto out = torch::empty_like(x);\n \n const int threads = 128;\n const int blocks = batch;\n \n // Calculate shared memory size\n // x_smem + x_norm + q + k + v + attn_scores + attn_out\n size_t shared_mem = SEQ_LEN * DIM * 5 * sizeof(float) + // x, x_norm, q, k, v\n NUM_HEADS * SEQ_LEN * SEQ_LEN * sizeof(float) + // attn_scores\n SEQ_LEN * DIM * sizeof(float); // attn_out\n \n // Use templates for fixed sizes\n if (seq_len == 32 && dim == 64 && num_heads == 4) {\n constexpr int SL = 32;\n constexpr int D = 64;\n constexpr int NH = 4;\n constexpr int HD = 16;\n size_t smem_needed = SL * D * 5 * sizeof(float) + NH * SL * SL * sizeof(float) + SL * D * sizeof(float);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"fused_attention_kernel\", ([&] {\n fused_attention_kernel<scalar_t, SL, D, NH, HD><<<blocks, threads, smem_needed>>>(\n x.data_ptr<scalar_t>(),\n norm_weight.data_ptr<scalar_t>(),\n norm_bias.data_ptr<scalar_t>(),\n qkv_weight.data_ptr<scalar_t>(),\n out_weight.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n batch\n );\n }));\n } else {\n // Fallback for other sizes (not optimized)\n AT_ERROR(\"Unsupported dimensions for fused attention kernel\");\n }\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor fused_attention_cuda(\n torch::Tensor x,\n torch::Tensor norm_weight,\n torch::Tensor norm_bias,\n torch::Tensor qkv_weight,\n torch::Tensor out_weight,\n int num_heads,\n int seq_len,\n int dim);\n\"\"\"\n\n# Load the extension\nattention_extension = load_inline(\n name='attention_extension_v2',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['fused_attention_cuda'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, dim, num_heads, height, width, dropout=0.0):\n super().__init__()\n self.height = height\n self.width = width\n self.num_heads = num_heads\n self.head_dim = dim // num_heads\n assert dim % num_heads == 0\n\n self.row_norm = nn.LayerNorm(dim)\n self.row_qkv = nn.Linear(dim, 3 * dim, bias=False)\n self.row_out = nn.Linear(dim, dim)\n\n self.col_norm = nn.LayerNorm(dim)\n self.col_qkv = nn.Linear(dim, 3 * dim, bias=False)\n self.col_out = nn.Linear(dim, dim)\n\n self.dropout = nn.Dropout(dropout)\n \n # Reference the extension\n self.fused_attention = attention_extension.fused_attention_cuda\n\n def _attn_fused(self, x, qkv_proj, out_proj, norm, seq_len):\n \"\"\"Attention using fused CUDA kernel.\"\"\"\n # This assumes no dropout for the fused kernel\n return self.fused_attention(\n x.contiguous(),\n norm.weight.contiguous(),\n norm.bias.contiguous(),\n qkv_proj.weight.transpose(0, 1).contiguous(), # [dim, 3*dim] -> [3*dim, dim]\n out_proj.weight.transpose(0, 1).contiguous(), # [dim, dim] -> [dim, dim]\n self.num_heads,\n seq_len,\n x.size(-1)\n )\n\n def _attn_ref(self, x, qkv_proj, out_proj, norm):\n \"\"\"Reference PyTorch attention for fallback.\"\"\"\n B_outer, L, D = x.shape\n h, d = self.num_heads, self.head_dim\n x_norm = norm(x)\n qkv = qkv_proj(x_norm).view(B_outer, L, 3, h, d)\n q, k, v = qkv.unbind(dim=2)\n q = q.transpose(1, 2)\n k = k.transpose(1, 2)\n v = v.transpose(1, 2)\n attn = (q @ k.transpose(-1, -2)) / math.sqrt(d)\n attn = F.softmax(attn, dim=-1)\n attn = self.dropout(attn)\n out = (attn @ v).transpose(1, 2).reshape(B_outer, L, D)\n return out_proj(out)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, H, W, D = x.shape\n \n # Check if we can use the fused kernel (dropout must be 0)\n use_fused = (self.dropout.p == 0.0 and D == 64 and self.num_heads == 4 and \n (W == 32 or H == 32))\n\n # Row-wise attention\n x_row = x.reshape(B * H, W, D)\n if use_fused and W == 32:\n row_out = self._attn_fused(x_row, self.row_qkv, self.row_out, self.row_norm, W)\n else:\n row_out = self._attn_ref(x_row, self.row_qkv, self.row_out, self.row_norm)\n x = x + row_out.reshape(B, H, W, D)\n\n # Column-wise attention\n x_col = x.permute(0, 2, 1, 3).reshape(B * W, H, D)\n if use_fused and H == 32:\n col_out = self._attn_fused(x_col, self.col_qkv, self.col_out, self.col_norm, H)\n x = x + col_out.reshape(B, W, H, D).permute(0, 2, 1, 3)\n else:\n col_out = self._attn_ref(x_col, self.col_qkv, self.col_out, self.col_norm)\n x = x + col_out.reshape(B, W, H, D).permute(0, 2, 1, 3)\n\n return x"
}