reasoning
Now I understand the GPU specs. Let me write a CUDA kernel for this scene graph relation update.
The key operations are:
1. LayerNorm on nodes
2. Q, K, V projections (linear layers without bias)
3. Rel MLP: Linear -> SiLU -> Linear
4. Attention computation with pairwise bias
5. Output projection
6. Residual add
For small batch sizes (B=4, N=24), we need to think about how to parallelize efficiently. The matrix multiplications are small:
- Q, K, V projections: (B*N, D) x (D, D) = (B*N, D)
- Attention: (B, H, N, head_dim) x (B, H, head_dim, N) = (B, H, N, N) - small!
- Rel MLP: (B, N, N, P) -> (B, N, N, P) -> (B, N, N, H)
- Output projection: (B*N, D) x (D, D) = (B*N, D)
The matrices are relatively small, so we need to be careful about kernel design. Let me think about a good approach:
1. For the projections (Q, K, V, out), we have small matrices. We could fuse these with other operations.
2. The attention computation involves small matrices (N=24), so we might benefit from shared memory tiling.
3. The rel_mlp operates on (B, N, N, P) and produces (B, H, N, N) bias for attention.
Let me start with a kernel that fuses all operations together. Given the small problem size, I'll use a block-per-(batch, head, query) approach for the attention computation.
Actually, let me think about this more carefully. The operations are:
1. LayerNorm: Can be done with a simple kernel
2. Q, K, V projections: These are matmuls with weights from the model
3. Rel MLP: Two matmuls with SiLU activation between them
4. Attention: Matmul + softmax + matmul with bias
5. Output projection: Mat
This is complex. Let me write a fused kernel that does:
- Load all weights from the model
- Compute layer norm
- Compute Q, K, V projections
- Compute rel_mlp for bias
- Compute attention
- Compute output projection
- Add residual
Actually, for this problem, since the matrices are small and there are many small operations, let me write a custom kernel that fuses the attention computation with the pairwise bias.
Let me start by writing a simpler approach - using PyTorch for the projections and only customizing the attention part. Then I can expand to fuse more operations.
Actually, let me think about a different approach. The key insight is that for small matrices, we can use one thread block per attention head and query, and have all threads in the block cooperate on the computation.
Let me design a fused kernel:
1. Thread blocks: (B, H, N) total, each block has WARP_SIZE threads
2. Each block computes the output for one query position
3. Within each block:
- Load Q for this query position
- Load K and V for all positions (shared memory)
- Compute rel_bias for all positions (compute on the fly or load)
- Compute attention scores, apply mask, softmax
- Compute weighted sum of V
- Store result
Wait, this approach still has many separate kernel launches for projections. Let me think about a more integrated approach.
Given the problem size (B=4, N=24, D=96, H=4), let me try a block-per-batch approach where each block handles all heads and all query positions within a batch. This would allow us to share the K, V values across query positions.
Actually, let me just start coding and iterate. I'll write a kernel that fuses the attention computation with the bias computation, and use PyTorch for the projections.
Hmm, but actually, for the MLP on pair_feats, that's a lot of computation (B*N*N = 4*24*24 = 2304 positions, each with P=16 features going through two linear layers). Let me see if we can fuse this efficiently.
Let me write a first version of the kernel:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n\n# Define the CUDA kernel\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <c10/cuda/CUDAException.h>\n\nconstexpr int WARP_SIZE = 32;\nconstexpr int MAX_THREADS = 256;\n\ntemplate <typename scalar_t>\n__device__ __forceinline__ scalar_t warpReduceSum(scalar_t val) {\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xffffffff, val, offset);\n }\n return val;\n}\n\ntemplate <typename scalar_t>\n__device__ __forceinline__ scalar_t blockReduceSum(scalar_t val, int tid, int blockDim) {\n static __shared__ scalar_t shared[32];\n int lane = tid % WARP_SIZE;\n int wid = tid / WARP_SIZE;\n \n val = warpReduceSum(val);\n \n if (lane == 0) {\n shared[wid] = val;\n }\n __syncthreads();\n \n if (wid == 0) {\n val = (tid < blockDim / WARP_SIZE) ? shared[lane] : scalar_t(0);\n val = warpReduceSum(val);\n }\n return val;\n}\n\n// LayerNorm kernel\ntemplate <typename scalar_t>\n__global__ void layer_norm_kernel(\n const scalar_t* __restrict__ input,\n const scalar_t* __restrict__ weight,\n const scalar_t* __restrict__ bias,\n scalar_t* __restrict__ output,\n int B, int N, int D, float eps) {\n \n int batch_idx = blockIdx.x / N;\n int node_idx = blockIdx.x % N;\n int tid = threadIdx.x;\n \n if (batch_idx >= B) return;\n \n const scalar_t* in_ptr = input + (batch_idx * N + node_idx) * D;\n scalar_t* out_ptr = output + (batch_idx * N + node_idx) * D;\n \n // Compute mean\n scalar_t sum = 0;\n for (int d = tid; d < D; d += blockDim.x) {\n sum += in_ptr[d];\n }\n scalar_t mean = blockReduceSum(sum, tid, blockDim.x);\n __shared__ scalar_t s_mean;\n if (tid == 0) s_mean = mean / D;\n __syncthreads();\n \n // Compute variance\n scalar_t var_sum = 0;\n for (int d = tid; d < D; d += blockDim.x) {\n scalar_t diff = in_ptr[d] - s_mean;\n var_sum += diff * diff;\n }\n scalar_t var = blockReduceSum(var_sum, tid, blockDim.x);\n __shared__ scalar_t s_rstd;\n if (tid == 0) s_rstd = rsqrtf(var / D + eps);\n __syncthreads();\n \n // Normalize and apply affine transform\n for (int d = tid; d < D; d += blockDim.x) {\n out_ptr[d] = (in_ptr[d] - s_mean) * s_rstd * weight[d] + bias[d];\n }\n}\n\n// Matrix multiplication kernel for small matrices (B*N, D) x (D, D)\ntemplate <typename scalar_t>\n__global__ void linear_kernel(\n const scalar_t* __restrict__ input,\n const scalar_t* __restrict__ weight,\n scalar_t* __restrict__ output,\n int M, int K, int N) {\n \n int row = blockIdx.x;\n int col = threadIdx.x;\n \n if (row >= M) return;\n \n scalar_t sum = 0;\n for (int k = 0; k < K; k++) {\n sum += input[row * K + k] * weight[k * N + col];\n }\n output[row * N + col] = sum;\n}\n\n// Rel MLP kernel: Linear -> SiLU -> Linear\ntemplate <typename scalar_t>\n__global__ void rel_mlp_kernel(\n const scalar_t* __restrict__ pair_feats,\n const scalar_t* __restrict__ weight1,\n const scalar_t* __restrict__ weight2,\n scalar_t* __restrict__ output,\n int B, int N, int P, int num_heads) {\n \n int idx = blockIdx.x;\n int tid = threadIdx.x;\n \n int b = idx / (N * N);\n int rem = idx % (N * N);\n int i = rem / N;\n int j = rem % N;\n \n if (b >= B) return;\n \n const scalar_t* feat_ptr = pair_feats + idx * P;\n \n // First linear: P -> P\n extern __shared__ char shared_mem[];\n scalar_t* intermediate = (scalar_t*)shared_mem;\n \n for (int p = tid; p < P; p += blockDim.x) {\n scalar_t sum = 0;\n for (int k = 0; k < P; k++) {\n sum += feat_ptr[k] * weight1[k * P + p];\n }\n // SiLU activation: x * sigmoid(x)\n scalar_t sigmoid_val = 1.0f / (1.0f + expf(-sum));\n intermediate[p] = sum * sigmoid_val;\n }\n __syncthreads();\n \n // Second linear: P -> num_heads\n for (int h = tid; h < num_heads; h += blockDim.x) {\n scalar_t sum = 0;\n for (int k = 0; k < P; k++) {\n sum += intermediate[k] * weight2[k * num_heads + h];\n }\n output[(b * num_heads + h) * N * N + i * N + j] = sum;\n }\n}\n\n// Attention kernel with bias\ntemplate <typename scalar_t>\n__global__ void attention_kernel(\n const scalar_t* __restrict__ q,\n const scalar_t* __restrict__ k,\n const scalar_t* __restrict__ v,\n const scalar_t* __restrict__ rel_bias,\n const float* __restrict__ edge_mask,\n scalar_t* __restrict__ ctx,\n int B, int N, int num_heads, int head_dim, float scale) {\n \n int b = blockIdx.x / num_heads;\n int h = blockIdx.x % num_heads;\n int q_idx = blockIdx.y;\n int tid = threadIdx.x;\n \n if (b >= B) return;\n \n // Shared memory for K, V across positions\n extern __shared__ char smem[];\n scalar_t* s_k = (scalar_t*)smem;\n scalar_t* s_v = s_k + N * head_dim;\n scalar_t* s_scores = s_v + N * head_dim;\n float* s_attn = (float*)(s_scores + N);\n \n const scalar_t* q_ptr = q + (b * num_heads + h) * N * head_dim + q_idx * head_dim;\n const scalar_t* k_ptr = k + (b * num_heads + h) * N * head_dim;\n const scalar_t* v_ptr = v + (b * num_heads + h) * N * head_dim;\n const scalar_t* bias_ptr = rel_bias + (b * num_heads + h) * N * N + q_idx * N;\n const float* mask_ptr = edge_mask + b * N * N + q_idx * N;\n scalar_t* out_ptr = ctx + (b * num_heads + h) * N * head_dim + q_idx * head_dim;\n \n // Load Q into registers\n scalar_t q_reg[32]; // Max head_dim\n for (int d = tid; d < head_dim; d += blockDim.x) {\n q_reg[d] = q_ptr[d];\n }\n \n // Load K and V into shared memory\n for (int pos = tid; pos < N; pos += blockDim.x) {\n for (int d = 0; d < head_dim; d++) {\n s_k[pos * head_dim + d] = k_ptr[pos * head_dim + d];\n s_v[pos * head_dim + d] = v_ptr[pos * head_dim + d];\n }\n }\n __syncthreads();\n \n // Compute attention scores\n for (int kv_idx = tid; kv_idx < N; kv_idx += blockDim.x) {\n scalar_t score = 0;\n for (int d = 0; d < head_dim; d++) {\n score += q_reg[d] * s_k[kv_idx * head_dim + d];\n }\n score = score * scale + bias_ptr[kv_idx];\n \n // Apply mask\n if (mask_ptr[kv_idx] < 0.5f) {\n s_scores[kv_idx] = -INFINITY;\n } else {\n s_scores[kv_idx] = score;\n }\n }\n __syncthreads();\n \n // Softmax reduction (find max for stability)\n scalar_t max_score = -INFINITY;\n for (int i = tid; i < N; i += blockDim.x) {\n max_score = max(max_score, s_scores[i]);\n }\n max_score = blockReduceSum(max_score, tid, blockDim.x);\n __shared__ scalar_t s_max;\n if (tid == 0) s_max = max_score;\n __syncthreads();\n \n // Compute exp(score - max) and sum\n scalar_t exp_sum = 0;\n for (int i = tid; i < N; i += blockDim.x) {\n if (s_scores[i] > -INFINITY) {\n s_attn[i] = expf(s_scores[i] - s_max);\n exp_sum += s_attn[i];\n } else {\n s_attn[i] = 0;\n }\n }\n exp_sum = blockReduceSum(exp_sum, tid, blockDim.x);\n __shared__ scalar_t s_sum;\n if (tid == 0) s_sum = exp_sum;\n __syncthreads();\n \n // Normalize attention weights\n for (int i = tid; i < N; i += blockDim.x) {\n s_attn[i] = s_attn[i] / s_sum;\n }\n __syncthreads();\n \n // Compute weighted sum of V\n for (int d = tid; d < head_dim; d += blockDim.x) {\n scalar_t val = 0;\n for (int i = 0; i < N; i++) {\n val += s_attn[i] * s_v[i * head_dim + d];\n }\n out_ptr[d] = val;\n }\n}\n\n// Output projection kernel\ntemplate <typename scalar_t>\n__global__ void out_proj_kernel(\n const scalar_t* __restrict__ ctx,\n const scalar_t* __restrict__ weight,\n scalar_t* __restrict__ output,\n int B, int N, int D) {\n \n int row = blockIdx.x;\n int col = threadIdx.x;\n \n if (row >= B * N) return;\n \n scalar_t sum = 0;\n for (int k = 0; k < D; k++) {\n sum += ctx[row * D + k] * weight[k * D + col];\n }\n output[row * D + col] = sum;\n}\n\n// Residual add kernel\ntemplate <typename scalar_t>\n__global__ void residual_add_kernel(\n const scalar_t* __restrict__ input,\n const scalar_t* __restrict__ proj_out,\n scalar_t* __restrict__ output,\n int size) {\n \n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < size) {\n output[idx] = input[idx] + proj_out[idx];\n }\n}\n\ntorch::Tensor forward_cuda(\n torch::Tensor nodes,\n torch::Tensor pair_feats,\n torch::Tensor edge_mask,\n torch::Tensor norm_weight,\n torch::Tensor norm_bias,\n torch::Tensor q_weight,\n torch::Tensor k_weight,\n torch::Tensor v_weight,\n torch::Tensor rel_weight1,\n torch::Tensor rel_weight2,\n torch::Tensor out_weight,\n int num_heads,\n float eps) {\n \n int B = nodes.size(0);\n int N = nodes.size(1);\n int D = nodes.size(2);\n int P = pair_feats.size(3);\n int head_dim = D / num_heads;\n float scale = 1.0f / sqrtf((float)head_dim);\n \n auto options = nodes.options();\n \n // Allocate intermediate tensors\n auto norm_out = torch::empty({B, N, D}, options);\n auto q = torch::empty({B, num_heads, N, head_dim}, options);\n auto k = torch::empty({B, num_heads, N, head_dim}, options);\n auto v = torch::empty({B, num_heads, N, head_dim}, options);\n auto rel_bias = torch::empty({B, num_heads, N, N}, options);\n auto ctx = torch::empty({B, num_heads, N, head_dim}, options);\n auto ctx_reshaped = torch::empty({B, N, D}, options);\n auto proj_out = torch::empty({B, N, D}, options);\n auto output = torch::empty({B, N, D}, options);\n \n int threads = min(D, MAX_THREADS);\n \n // LayerNorm\n AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), \"layer_norm\", ([&] {\n layer_norm_kernel<scalar_t><<<B * N, threads>>>(\n nodes.data_ptr<scalar_t>(),\n norm_weight.data_ptr<scalar_t>(),\n norm_bias.data_ptr<scalar_t>(),\n norm_out.data_ptr<scalar_t>(),\n B, N, D, eps);\n }));\n \n // Q, K, V projections\n threads = min(D, MAX_THREADS);\n AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), \"q_proj\", ([&] {\n linear_kernel<scalar_t><<<B * N, threads>>>(\n norm_out.data_ptr<scalar_t>(),\n q_weight.data_ptr<scalar_t>(),\n ctx_reshaped.data_ptr<scalar_t>(), // reuse buffer\n B * N, D, D);\n }));\n \n // Reshape to (B, num_heads, N, head_dim)\n auto q_temp = ctx_reshaped.view({B, N, num_heads, head_dim}).permute({0, 2, 1, 3});\n q.copy_(q_temp);\n \n AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), \"k_proj\", ([&] {\n linear_kernel<scalar_t><<<B * N, threads>>>(\n norm_out.data_ptr<scalar_t>(),\n k_weight.data_ptr<scalar_t>(),\n ctx_reshaped.data_ptr<scalar_t>(),\n B * N, D, D);\n }));\n auto k_temp = ctx_reshaped.view({B, N, num_heads, head_dim}).permute({0, 2, 1, 3});\n k.copy_(k_temp);\n \n AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), \"v_proj\", ([&] {\n linear_kernel<scalar_t><<<B * N, threads>>>(\n norm_out.data_ptr<scalar_t>(),\n v_weight.data_ptr<scalar_t>(),\n ctx_reshaped.data_ptr<scalar_t>(),\n B * N, D, D);\n }));\n auto v_temp = ctx_reshaped.view({B, N, num_heads, head_dim}).permute({0, 2, 1, 3});\n v.copy_(v_temp);\n \n // Rel MLP\n threads = min(P, MAX_THREADS);\n size_t shared_mem_size = P * sizeof(float) * 2; // intermediate values\n AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), \"rel_mlp\", ([&] {\n rel_mlp_kernel<scalar_t><<<B * N * N, threads, shared_mem_size>>>(\n pair_feats.data_ptr<scalar_t>(),\n rel_weight1.data_ptr<scalar_t>(),\n rel_weight2.data_ptr<scalar_t>(),\n rel_bias.data_ptr<scalar_t>(),\n B, N, P, num_heads);\n }));\n \n // Attention\n threads = min(N, MAX_THREADS);\n size_t attn_shared_mem = N * head_dim * 2 * sizeof(float) + N * sizeof(float) + N * sizeof(float);\n dim3 grid(B * num_heads, N);\n dim3 block(threads);\n \n AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), \"attention\", ([&] {\n attention_kernel<scalar_t><<<grid, block, attn_shared_mem>>>(\n q.data_ptr<scalar_t>(),\n k.data_ptr<scalar_t>(),\n v.data_ptr<scalar_t>(),\n rel_bias.data_ptr<scalar_t>(),\n edge_mask.data_ptr<float>(),\n ctx.data_ptr<scalar_t>(),\n B, N, num_heads, head_dim, scale);\n }));\n \n // Reshape context\n auto ctx_temp = ctx.permute({0, 2, 1, 3}).reshape({B, N, D});\n ctx_reshaped.copy_(ctx_temp);\n \n // Output projection\n threads = min(D, MAX_THREADS);\n AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), \"out_proj\", ([&] {\n out_proj_kernel<scalar_t><<<B * N, threads>>>(\n ctx_reshaped.data_ptr<scalar_t>(),\n out_weight.data_ptr<scalar_t>(),\n proj_out.data_ptr<scalar_t>(),\n B, N, D);\n }));\n \n // Residual add\n int size = B * N * D;\n threads = MAX_THREADS;\n int blocks = (size + threads - 1) / threads;\n AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), \"residual_add\", ([&] {\n residual_add_kernel<scalar_t><<<blocks, threads>>>(\n nodes.data_ptr<scalar_t>(),\n proj_out.data_ptr<scalar_t>(),\n output.data_ptr<scalar_t>(),\n size);\n }));\n \n return output;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor forward_cuda(\n torch::Tensor nodes,\n torch::Tensor pair_feats,\n torch::Tensor edge_mask,\n torch::Tensor norm_weight,\n torch::Tensor norm_bias,\n torch::Tensor q_weight,\n torch::Tensor k_weight,\n torch::Tensor v_weight,\n torch::Tensor rel_weight1,\n torch::Tensor rel_weight2,\n torch::Tensor out_weight,\n int num_heads,\n float eps);\n\"\"\"\n\n# Load the extension\ngf_vla_ops = load_inline(\n name='gf_vla_ops',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['forward_cuda'],\n extra_cuda_cflags=['-O3'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, dim: int, pair_dim: int, num_heads: int):\n super().__init__()\n assert dim % num_heads == 0\n self.dim = dim\n self.pair_dim = pair_dim\n self.num_heads = num_heads\n self.head_dim = dim // num_heads\n self.scale = 1.0 / math.sqrt(self.head_dim)\n\n self.norm = nn.LayerNorm(dim)\n self.q_proj = nn.Linear(dim, dim, bias=False)\n self.k_proj = nn.Linear(dim, dim, bias=False)\n self.v_proj = nn.Linear(dim, dim, bias=False)\n self.rel_mlp = nn.Sequential(\n nn.Linear(pair_dim, pair_dim),\n nn.SiLU(),\n nn.Linear(pair_dim, num_heads),\n )\n self.out_proj = nn.Linear(dim, dim, bias=False)\n\n def forward(\n self,\n nodes: torch.Tensor,\n pair_feats: torch.Tensor,\n edge_mask: torch.Tensor,\n ) -> torch.Tensor:\n # Ensure inputs are contiguous\n nodes = nodes.contiguous()\n pair_feats = pair_feats.contiguous()\n edge_mask = edge_mask.contiguous()\n \n # Get weights from the model\n norm_weight = self.norm.weight.contiguous()\n norm_bias = self.norm.bias.contiguous()\n q_weight = self.q_proj.weight.contiguous()\n k_weight = self.k_proj.weight.contiguous()\n v_weight = self.v_proj.weight.contiguous()\n rel_weight1 = self.rel_mlp[0].weight.contiguous()\n rel_weight2 = self.rel_mlp[2].weight.contiguous()\n out_weight = self.out_proj.weight.contiguous()\n \n # Call the CUDA kernel\n return gf_vla_ops.forward_cuda(\n nodes, pair_feats, edge_mask,\n norm_weight, norm_bias,\n q_weight, k_weight, v_weight,\n rel_weight1, rel_weight2, out_weight,\n self.num_heads, 1e-5\n )"
}