We need to write a custom CUDA kernel that computes the forward pass of the SE(3)-invariant linear layer faster than PyTorch. The layer: - Input: scalar_feat (B, N, scalar_in), vector_feat (B, N, vector_in, 3) - Compute vec_norm = vector_feat.norm(dim=-1) -> shape (B, N, vector_in) - combined = cat([scalar_feat, vec_norm], dim=-1) -> shape (B, N, scalar_in + vector_in) - s_out = combined @ W_s^T + b_s where W_s,b_s from scalar_linear (nn.Linear) - v_out = einsum("oi,bnid->bnod", vector_weight, vector_feat) -> shape (B, N, vector_out, 3). Actually vector_weight shape (vector_out, vector_in). So v_out[b,n,o,i3?] = sum_i vector_weight[o,i] * vector_feat[b,n,i,3]. So it's just a linear combination of input channels into output channels per node, per coordinate. - gate = sigmoid(combined @ W_g^T + b_g) where W_g,b_g from vector_gate (nn.Linear) -> shape (B, N, vector_out). Then unsqueeze(-1) -> (B, N, vector_out, 1), multiply v_out. - Return s_out, v_out. Optimization opportunities: - Fusing all operations into one kernel: compute vec_norms, combine, compute scalar linear and gate linear, compute vector linear, apply sigmoid and multiply. - We can also pre-load weights to shared memory or registers. - For linear layers, we can use standard GEMM but fused with norms and gates. However, writing a fully optimized GEMM is huge. But the sizes are relatively small: scalar_in=64, scalar_out=64, vector_in=16, vector_out=16. Combined size = 80. So scalar_linear is 80->64, vector_gate is 80->16. vector_weight is 16x16. The batch size is B=8, N=128, total nodes = 1024. This is quite small. PyTorch reference will probably launch multiple kernels: norm, cat, linear (GEMM), einsum (maybe bmm or permute+matmul), linear for gate, sigmoid, mul. The overhead of multiple kernels might be significant. A fused CUDA kernel per node (or per thread) should be efficient. Since total nodes = 1024, each node does: - Load vector_feat of shape (vector_in, 3) = 16*3=48 floats. - Compute squared norms per vector_in channel: sum over 3 dims. 16 norms. - scalar_in = 64. So combined is 80 floats. - Compute s_out = W_s * combined + b_s. W_s is (64,80). This is a small matrix-vector multiply per node. 64*80=5120 multiplications per node. - Compute gate = W_g * combined + b_g. W_g is (16,80). 1280 multiplications per node. - Compute v_out = W_v * vector_feat. W_v is (16,16). vector_feat is (16,3). So v_out (16,3) = W_v (16,16) * vector_feat (16,3). This is 16*16*3 = 768 multiplications per node. - Apply sigmoid to gate (16), multiply each of 16 output channels by gate value across 3 coordinates: 16*3 = 48 multiplications. - Write s_out (64), v_out (16*3=48). Total FLOPs per node: ~5120+1280+768+something ~7200 FLOPs. Total nodes: 1024. Total FLOPs: ~7.4 MFLOPs. Very tiny. Memory traffic: read vector_feat (48 floats), read scalar_feat (64 floats), write s_out (64), write v_out (48). Plus weights (read once, maybe cached). Total per node: ~224 floats = ~900 bytes. Total memory: ~900 KB. This is tiny, bandwidth bound if not fused, but compute is also tiny. Fusing into one kernel reduces kernel launch overhead and memory roundtrips for intermediate tensors (combined, vec_norm, gate, v_out before gating). Since the problem size is small, kernel launch overhead is likely dominant. A single kernel should be much faster. Implementation plan: Write a CUDA kernel where each thread handles one node (b,n). So grid of B*N threads, maybe grouped into blocks of e.g. 128 or 256 threads. Within each thread: - Load scalar_feat[scalar_in] into registers. - Load vector_feat[vector_in][3] into registers or shared memory. Since vector_in=16, 3 dims, 48 floats. That's a lot for registers but okay per thread (48 registers). Actually each thread can hold 48 floats, plus weights, etc. - Compute vec_norm_sq or just vec_norm (sqrt). The reference uses .norm(dim=-1) which computes sqrt(sum(x^2)). We need sqrt. But wait: The reference uses `vec_norm = vector_feat.norm(dim=-1)` which is L2 norm with sqrt. Then `combined = torch.cat([scalar_feat, vec_norm], dim=-1)`. So we need sqrt. - Assemble combined vector of length scalar_in + vector_in = 80. We can keep scalar_feat values in registers, then append vec_norm values. So we don't need to allocate a shared array; we can just use the scalar_feat values and then norms when doing the matmuls. - For scalar_linear: W_s is shape (scalar_out, scalar_in+vector_in). We'll store W_s in a 1D array of size scalar_out * combined_size. Same for b_s. - For vector_gate: W_g is shape (vector_out, combined_size). - For vector_weight: W_v is shape (vector_out, vector_in). Since these sizes are small and fixed at compile time (or known at runtime), we can pass them as template parameters for efficiency, but it's okay to pass as int arguments. However, the kernel should handle arbitrary sizes as per constructor args, but the benchmark uses the fixed sizes. We can either write a generic kernel or use template instantiation. For simplicity, we can write a generic kernel with loops. But we must be careful about register usage. If each thread does many loops, it may spill. However, the sizes are small. Let's think about block design. Option 1: One thread per node. Total nodes=1024. Launch 1024 threads. Each thread does ~7200 FLOPs. At 20 TFLOPS, that's 0.36 microseconds, but memory latency will dominate. 1024 threads may not saturate the GPU. Better to use more threads. Option 2: One warp per node. 32 threads cooperate. Then total warps = 1024/32 = 32 warps, way too few. But each warp can work on one node, doing the matmuls cooperatively. 32 warps across ~80 SMs means very underutilized. But kernel launch overhead is the main bottleneck. Still, with 1024 threads, occupancy may be low but it's okay for tiny problems? The GPU is likely an A100 or H100. 1024 threads is one wave on one SM. The kernel will be extremely fast. Actually the reference PyTorch might also be fast. But fused kernel can be faster due to fewer memory roundtrips and kernel launches. Let's design with one thread per node but unroll loops. Or maybe one block per node with multiple threads? With 1024 nodes, blocks of size 128 => 8 blocks. That's only 8 blocks. That's too few. We need more parallelism. Better: each thread handles one element of the output? But then we need communication for norms and gating. Alternative: block size = e.g. 256 threads, grid = (B*N*something). Actually we can handle multiple nodes per block. For example, block size = 256 threads, each thread handles one node, but a block processes 256 nodes. Then we need 4 blocks (since 1024 nodes). That's 4 blocks. Still tiny. If we process multiple elements per thread, e.g., each thread handles K nodes, then total threads = 1024/K. We want enough blocks to fill SMs. K=1 => 1024 threads. Actually we can also assign each thread to handle one output channel for one node. For scalar_out=64, vector_out*3=48, total outputs per node = 112. If each thread computes one output, we need 1024*112 = 114,688 threads. That's a decent amount. We can have each thread compute one output element of one node. That yields ~114k threads, which is enough to occupy the GPU. But then threads within a block need to cooperate to compute combined (shared) and gate (shared). However, combined depends on vector norms, which require reading vector_feat. If each thread reads only part of vector_feat, it gets complicated. A simpler design: assign each warp (32 threads) to one node. Then total warps = 1024. That's 32 blocks of 1024 threads? Wait, if block size = 128 threads (4 warps), and each warp handles one node, then one block handles 4 nodes. To handle 1024 nodes, we need 256 blocks. That's plenty! And each warp cooperates on the computations: compute norms cooperatively, compute linear layers using warp-level reductions or just partitioning the rows. Each warp (32 threads) can compute the scalar linear (64 outputs) by having each thread compute 2 outputs (64/32=2). Similarly, vector transform: 16 outputs * 3 dims = 48 elements, each thread can compute 1 or 2. Gate: 16 outputs, each thread computes 0.5? Better to have each thread compute a subset. But warp-level programming is a bit more code, but not too bad. However, implementing a general warp parallel GEMM for arbitrary sizes is nontrivial but manageable for these tiny sizes. Even simpler: one thread per node, but process multiple nodes per thread in a grid-stride loop. So each block has e.g. 256 threads, grid = enough blocks to get high occupancy. We can compute grid size to cover all nodes. Each thread loops over nodes with stride = total_threads. So if we launch e.g. 128 blocks of 256 threads = 32768 threads, each thread handles 1024/32768 = 0.031 nodes. In practice, many threads will do nothing, but the active threads (1024) will each handle one node. The advantage is that we launch many threads, but only 1024 do useful work. This doesn't help; it's the same. If we want to increase thread count to >1024, we need each node to be processed by multiple threads. So let's do one node per warp (or per block) but use warp-level parallelism to speed up each node. But because the problem is tiny, the actual speedup might come from fusing and keeping everything in registers/shared memory, reducing memory traffic. Launching 1024 threads may be enough. Let's think about PyTorch reference: it will do the following: 1. vector_feat.norm(dim=-1): launch a kernel, read vector_feat (B,N,Vin,3), write vec_norm (B,N,Vin). 1024*16*3=49152 reads, 1024*16 writes. 2. cat: launch a kernel? Actually cat may be fused into subsequent linear if PyTorch's fuser catches it? Unlikely. 3. scalar_linear: calls cuBLAS or a custom linear kernel. Input (1024,80) @ (80,64). This is a tiny GEMM, likely handled by a custom kernel or cuBLAS. cuBLAS might have overhead. 4. vector einsum: (1024,16,3) transformed by (16,16). This is like a batched matrix multiplication or grouped GEMM. PyTorch might call a kernel. 5. vector_gate: similar (1024,80) @ (80,16). 6. sigmoid: elementwise kernel. 7. mul: elementwise kernel. Total kernels: ~7-8. Each has launch overhead and global memory roundtrip. A fused kernel that does everything in one go should reduce latency significantly. Let's write a kernel where each block processes multiple nodes, and each thread processes one node. We can set block size to e.g. 128 or 256. Grid = (B*N + block_size -1)/block_size. So 1024/256 = 4 blocks. That is indeed few blocks. But the kernel will do a lot of compute per thread (register based), so the GPU occupancy is low but maybe still faster than reference because we avoid memory roundtrips. However, on modern GPUs, 4 blocks may run on 4 SMs, which is fine for a tiny problem. The latency will be dominated by memory reads (coalesced). Since the kernel is heavily memory-bound but fuses loads, it should be fine. But there's a nuance: if we do one thread per node, each thread reads scalar_feat (64 floats) and vector_feat (48 floats). With 256 threads per block, each block handles 256 nodes. The memory reads for scalar_feat and vector_feat will be coalesced if adjacent threads access adjacent nodes. Since scalar_feat is (B,N,scalar_in), indexing `scalar_feat[b,n,:]`, adjacent n are contiguous in memory if we index by linear index over (B,N). Let's map `idx = b * N + n`. Then for a given channel c, `scalar_feat[b,n,c]` is at offset `(b*N + n)*scalar_in + c`. So adjacent threads (adjacent idx) access the same channel c but offset by scalar_in. This is strided by scalar_in=64 floats = 256 bytes. For coalesced access, adjacent threads should access adjacent memory. Here, thread idx and idx+1 access addresses separated by scalar_in floats. That's contiguous if scalar_in is small? Actually, memory is arranged as: for node 0, channels 0..63, then node 1, channels 0..63, etc. So `scalar_feat[idx, c]` has address `base + (idx * scalar_in + c) * sizeof(float)`. Adjacent threads (idx, idx+1) accessing the same c have addresses separated by `scalar_in * sizeof(float)` = 256 bytes. That's a stride. For 128-bit loads (4 floats), a warp of 32 threads would load 32*256 = 8192 bytes apart, which is not contiguous. The coalescing unit on A100 is 128 bytes per warp? Actually, each warp issues 128-byte transactions. If threads access addresses spaced 256 bytes apart, each thread touches a different 128-byte segment, resulting in 32 transactions of 128 bytes = 4096 bytes loaded, but only 128 bytes actually used (4 bytes per thread). That's terrible bandwidth utilization. To improve memory coalescing, we should have adjacent threads load adjacent floats. That means each thread should handle a specific channel, and we use multiple threads per node. So one node is processed by a block of threads (or warp) that cooperatively load data. Yes, this is important. For a memory-bound tiny kernel, coalescing matters. Let's design a block-per-node approach. Since there are 1024 nodes, we need 1024 blocks. Block size can be e.g. 64 or 128 threads. 1024 blocks is enough to fill the GPU (A100 has 108 SMs, each can run multiple blocks). Each block of 64 threads processes one node. 1024 blocks * 64 threads = 65536 threads. That's a good amount. But wait, 1024 blocks of 64 threads each might have low occupancy if each block uses a lot of registers/shared memory. But the workload per block is tiny: one node. 64 threads collaborating on ~112 outputs and 80 inputs. We can use shared memory to store inputs and intermediate results: - scalar_feat: 64 floats - vector_feat: 16*3 = 48 floats - combined: 80 floats - vec_norm: 16 floats - s_out: 64 floats - v_out: 16*3 = 48 floats - gate: 16 floats Total shared memory per block: (64+48+80+16+64+48+16) * 4 bytes = 336 * 4 = 1344 bytes. Very small. But we can also keep everything in registers and just have each thread compute part of the outputs. Let's think about the most natural parallelization for a block of T threads processing one node. Option: T = 64 threads. - 16 threads compute vec_norms (each handles one of 16 vector channels, reading 3 floats). - For scalar_linear (64 outputs) and vector_gate (16 outputs), we have 80 outputs total? Actually 64+16=80 outputs from combined. We can assign each of the 64 threads to compute one scalar_out and one vector_gate output. But vector_gate only has 16 outputs; we can assign 16 threads to compute both one scalar_out and one gate_out, and the rest 48 threads compute two scalar_outs each? Hmm. Better: use 64 threads. - Thread t (0..63) computes scalar_out[t] = dot(W_s[t, :], combined) + b_s[t]. (Assume scalar_out=64). This covers all scalar outputs. - For gate (16 outputs), threads 0..15 compute gate[t] = dot(W_g[t, :], combined) + b_g[t]. - For vector transform (vector_out=16, 3 dims = 48 outputs), threads 0..47 compute one coordinate each: v_out[t/3][t%3] = sum_i W_v[t/3][i] * vector_feat[i][t%3]. Or we can assign 64 threads: threads 0..47 compute one of 48, threads 48..63 are idle for this step. - Then gating: each gate value multiplies 3 coordinates. Threads 0..47 can do that. So within a block, we need to: 1. Load scalar_feat (64 floats) and vector_feat (48 floats) from global memory into shared memory or have each thread load its needed elements. - For scalar_feat: we need all 64 values for combined. We can have thread 0..63 each load one scalar_feat float. Let's store in shared memory `s_scalar[64]`. - For vector_feat: we need all 48 values. We can have thread 0..47 each load one coordinate. But vector_feat shape is (Vin,3). Let's store in `s_vector[48]` (or `s_vector[Vin][3]`). 2. Compute vec_norm: - Each thread t (0..15) reads s_vector[t][0..2], computes norm, stores in `s_combined[scalar_in + t]`. And maybe also `s_norm[t]`. - Threads 0..63 also copy scalar_feat into s_combined[0..scalar_in-1] (each thread copies its own). - __syncthreads(). 3. Compute scalar_linear and vector_gate: - Thread t (0..63) computes scalar_out[t] = dot(W_s[t*combined_size : (t+1)*combined_size], s_combined) + b_s[t]. Since combined_size = 80. That's a dot product of length 80. We can unroll or loop. With 64 threads, each doing a loop of 80, it's 5120 iterations total. Fine. - Threads 0..15 also compute gate[t] = dot(W_g[t*combined_size : (t+1)*combined_size], s_combined) + b_g[t]. After computing, apply sigmoid. - __syncthreads() maybe not needed if we store scalar_out in shared and gate in shared. 4. Compute vector linear: - Thread t (0..47) computes v_out[t] = sum_{i=0}^{Vin-1} W_v[(t/3)*Vin + i] * s_vector[i*3 + (t%3)]. - That's a loop of Vin=16 per thread. 48*16=768 iterations. - __syncthreads(). 5. Apply gate: - Thread t (0..47) reads gate[t/3] and multiplies v_out[t]. - Write v_out to global memory. - Thread t (0..63) writes scalar_out[t]. This is quite clean. Memory coalescing for loads of scalar_feat and vector_feat: - Since each block handles one node, reading scalar_feat for that node is not coalesced across threads within the block because all threads are reading from the same node? Wait, if each block handles one node, then threads in the block read different channels of the same node. In memory, channels are contiguous. So thread 0 reads channel 0, thread 1 reads channel 1, etc. So adjacent threads read adjacent floats. This is perfectly coalesced! Excellent. - Similarly for vector_feat: thread 0 reads first float of vector_feat for that node, thread 1 reads next float. Since vector_feat for a node is stored as (Vin,3) = 48 contiguous floats, reading them thread by thread is coalesced. - For weights: each thread reads a row of W_s. W_s is stored as (scalar_out, combined_size). Row t is at offset t*combined_size. Since each thread reads its own row, the accesses are strided by combined_size (80 floats) across threads. That's not coalesced. However, weights are read-only and can be cached in L1/L2. Since the total weight size is tiny (64*80 + 16*80 + 16*16 = 5120+1280+256 = 6656 floats = ~26KB), they will likely stay in L2 cache across all blocks. So coalescing for weights is less critical. But we could store weights in shared memory? Total weights = 6656 floats * 4 bytes = 26.6KB. Shared memory per SM is 164KB (A100). If each block loads all weights into shared memory, it's a waste because we're loading them for each block (1024 blocks), which is 26.6KB * 1024 = 27MB of memory traffic, which is still fine? Actually 27MB is fine. But each block would consume 26KB of shared memory, limiting occupancy to maybe 4-6 blocks per SM. But since we have 1024 blocks, and each block does little work, occupancy might not matter. However, loading weights from global memory to shared memory may not be beneficial because they are already cached. We can just read from global memory directly. Better: keep weights in constant memory? No, too large for constant cache. But L1/L2 will cache them nicely because all blocks read the same weights. Actually on A100, L2 cache is 40MB. 27KB is trivial. So we can just read weights from global memory and rely on cache. But we can also store weights in shared memory to exploit broadcast? Actually all threads read different rows, so shared memory doesn't help with broadcast. But if we transpose W_s so that columns are contiguous for coalescing? Let's think: if W_s is column-major (combined_size, scalar_out), then thread t reads column t, and adjacent threads read adjacent columns. But each dot product requires reading the entire column? No, dot product reads row t. If stored column-major, reading row t is strided by scalar_out. That's also strided. We can keep row-major and accept the strided reads, since it's cached. Alternatively, we can compute the dot product by having threads within a warp collaborate: each thread reads a chunk of combined and a chunk of W_s, accumulate partial sums, then warp-reduce. But given the small size (80), a simple per-thread loop is fine. However, we should be careful about register pressure: each thread loops over 80 elements for scalar, reads a weight row of 80, accumulates. That uses maybe a couple of registers for the loop. That's fine. Gate computation: threads 0..15 compute dot product of length 80. They read W_g. Same pattern. Vector linear: thread t computes for row o = t/3, coord d = t%3: sum_{i=0}^{Vin-1} W_v[o*Vin + i] * s_vector[i*3 + d]. So for a given coordinate d, the access to s_vector is strided by 3 (i*3+d). That's fine (shared memory). The access to W_v is row o, contiguous across i (since W_v is row-major (Vout, Vin)). So for a given thread, W_v reads are contiguous. Good. Across threads 0..47, W_v rows are accessed: threads 0,3,6,... (same o=0) read the same W_v row. That causes broadcast. Threads 1,4,7,... read row 1, etc. There is no bank conflict if using constant memory or L1. In shared memory, if we load W_v into shared memory, multiple threads reading the same row might have conflicts if they read the same bank? Actually if stored as `float W_v[Vout][Vin]`, threads reading the same row access the same banks sequentially or broadcast (depending on compute capability). On A100, broadcast is supported. So fine. But we don't need to put W_v in shared memory. Let's design the kernel signature: ```cpp torch::Tensor se3layer_forward( torch::Tensor scalar_feat, // (B, N, scalar_in) torch::Tensor vector_feat, // (B, N, vector_in, 3) torch::Tensor scalar_weight, // (scalar_out, scalar_in+vector_in) torch::Tensor scalar_bias, // (scalar_out) torch::Tensor vector_weight, // (vector_out, vector_in) torch::Tensor gate_weight, // (vector_out, scalar_in+vector_in) torch::Tensor gate_bias // (vector_out) ) -> std::vector<torch::Tensor> { s_out, v_out } ``` But we need to handle `scalar_in`, `scalar_out`, `vector_in`, `vector_out` dynamically. In CUDA kernel: ```cpp template <typename scalar_t> __global__ void se3kernel( const scalar_t* __restrict__ scalar_feat, const scalar_t* __restrict__ vector_feat, const scalar_t* __restrict__ scalar_weight, const scalar_t* __restrict__ scalar_bias, const scalar_t* __restrict__ vector_weight, const scalar_t* __restrict__ gate_weight, const scalar_t* __restrict__ gate_bias, scalar_t* __restrict__ scalar_out, scalar_t* __restrict__ vector_out, int B, int N, int scalar_in, int scalar_out, int vector_in, int vector_out) { // linear node index int node_idx = blockIdx.x; // one block per node int tid = threadIdx.x; int num_threads = blockDim.x; // e.g., 64 if (node_idx >= B * N) return; int b = node_idx / N; int n = node_idx % N; int combined_size = scalar_in + vector_in; // Shared memory extern __shared__ char smem[]; scalar_t* s_scalar = (scalar_t*)smem; // size scalar_in scalar_t* s_vector = s_scalar + scalar_in; // size vector_in * 3 scalar_t* s_combined = s_vector + vector_in * 3; // size combined_size scalar_t* s_scalar_out = s_combined + combined_size; // size scalar_out scalar_t* s_gate = s_scalar_out + scalar_out; // size vector_out scalar_t* s_vector_out = s_gate + vector_out; // size vector_out * 3 // Load scalar_feat for (int i = tid; i < scalar_in; i += num_threads) { s_scalar[i] = scalar_feat[((b * N + n) * scalar_in) + i]; } // Load vector_feat int vector_feat_size = vector_in * 3; for (int i = tid; i < vector_feat_size; i += num_threads) { s_vector[i] = vector_feat[((b * N + n) * vector_feat_size) + i]; } __syncthreads(); // Build combined: scalar part + vector norms // Copy scalar part for (int i = tid; i < scalar_in; i += num_threads) { s_combined[i] = s_scalar[i]; } // Compute vector norms // We'll have each thread handle some vector_in channels for (int i = tid; i < vector_in; i += num_threads) { scalar_t x = s_vector[i * 3 + 0]; scalar_t y = s_vector[i * 3 + 1]; scalar_t z = s_vector[i * 3 + 2]; s_combined[scalar_in + i] = sqrtf(x*x + y*y + z*z); } __syncthreads(); // Scalar linear for (int o = tid; o < scalar_out; o += num_threads) { scalar_t accum = 0; const scalar_t* w_row = scalar_weight + o * combined_size; for (int i = 0; i < combined_size; ++i) { accum += w_row[i] * s_combined[i]; } s_scalar_out[o] = accum + scalar_bias[o]; } // Vector gate for (int o = tid; o < vector_out; o += num_threads) { scalar_t accum = 0; const scalar_t* w_row = gate_weight + o * combined_size; for (int i = 0; i < combined_size; ++i) { accum += w_row[i] * s_combined[i]; } scalar_t g = accum + gate_bias[o]; // sigmoid g = 1.0f / (1.0f + expf(-g)); s_gate[o] = g; } __syncthreads(); // Vector linear int v_out_size = vector_out * 3; for (int idx = tid; idx < v_out_size; idx += num_threads) { int o = idx / 3; int d = idx % 3; scalar_t accum = 0; const scalar_t* w_row = vector_weight + o * vector_in; for (int i = 0; i < vector_in; ++i) { accum += w_row[i] * s_vector[i * 3 + d]; } s_vector_out[idx] = accum; } __syncthreads(); // Apply gate for (int idx = tid; idx < v_out_size; idx += num_threads) { int o = idx / 3; s_vector_out[idx] = s_vector_out[idx] * s_gate[o]; } __syncthreads(); // Write outputs for (int i = tid; i < scalar_out; i += num_threads) { scalar_out[((b * N + n) * scalar_out) + i] = s_scalar_out[i]; } for (int i = tid; i < v_out_size; i += num_threads) { vector_out[((b * N + n) * v_out_size) + i] = s_vector_out[i]; } } ``` This should work. But we need to be careful about shared memory size: we can compute it as `(scalar_in + vector_in*3 + combined_size + scalar_out + vector_out + vector_out*3) * sizeof(scalar_t)`. For the given sizes: 64 + 48 + 80 + 64 + 16 + 48 = 320 floats = 1280 bytes. Very small. But is a block size of 64 optimal? With 64 threads, each block has 2 warps. 1024 blocks total. On A100, 108 SMs, each can run up to 32 warps concurrently (2048 threads). So each SM can run up to 1024 blocks? No, blocks per SM limit is 32 (or 16). So total concurrent blocks = 108*32 = 3456. We have 1024 blocks, which is less than 3456, so all blocks can be in flight. But each block has only 2 warps, so total warps = 2048 warps. Occupancy per SM: 2048 warps / 108 SMs = ~19 warps per SM. The A100 has 64 warps per SM capacity. So occupancy is ~30%. Not great, but okay. We can increase block size to 128 or 256 to get higher occupancy? Actually if block size = 256 threads (8 warps), total blocks = 1024. Wait, no: if block size = 256, then we have 1024 blocks? No, one block per node = 1024 blocks regardless of block size. Block size just changes threads per block. Total threads = 1024 * block_size. If block_size=256, total threads = 262144. That's a lot. But each block does redundant work? No, each block handles one node, with 256 threads cooperating. For a node, there are only ~320 elements of work. So many threads will be idle (since loops stride by blockDim.x). With block_size=256, many threads won't do anything because the loop bounds are small (e.g., scalar_in=64, so only first 64 threads load, rest idle). That's okay, but it wastes thread slots. Better block size might be 64 or 128. Let's stick with 64 or 128. Let's try 128. Then total warps = 1024 * 4 = 4096 warps. That's better occupancy. But many threads idle in loops with small bounds. However, some loops are larger: scalar_out=64 but block_size=128 means 64 threads do work, 64 idle. Vector_out*3=48, 48 threads work, 80 idle. Not great. What if we assign each thread to handle multiple nodes? For example, block size = 256, grid = (B*N + block_size -1)/block_size = 4 blocks. Then each thread handles one node in a grid-stride loop? Actually blockIdx.x*blockDim.x + threadIdx.x gives linear thread id. Then we can have each thread handle one node if total threads >= B*N. If block size=256 and grid=4, total threads=1024, exactly B*N. Each thread handles exactly one node. Then we have only 4 blocks, but each block has 256 threads. Total warps = 32. That's extremely low. Alternatively, block size = 128, grid = 8 blocks. Total threads = 1024. Each thread handles one node. Total warps = 16. Even worse. So to get enough blocks, we need either one block per node (grid = B*N) or multiple blocks per node. One block per node with block size 64 gives 1024 blocks and 2048 warps. That's decent. Could we do multiple blocks per node? For example, split scalar_out across blocks? But then we'd need to communicate combined and vector data, which is complicated. Could we do one warp per node (block size = 32)? Then total warps = 1024. That's half of previous. Block size 64 seems better. Could we do 2 nodes per block? For example, block size = 128, each block handles 2 nodes. Then grid = 512 blocks. Total warps = 512 * 4 = 2048 warps (same as 1024 blocks of 64). But each node gets 64 threads. This might be better because it reduces block count but still enough. Actually 512 blocks is also fine. And we can process 2 nodes per block: first 64 threads handle node 0, next 64 handle node 1. This introduces some divergence? Not really if we branch on `threadIdx.x < 64`. But we need separate shared memory arrays for each node, or offset indexing. Simpler: one block per node, 64 threads. That's clean. Let's calculate shared memory for one block: 320 floats = 1280 bytes. With 1024 blocks active, total shared memory allocated across all resident blocks: e.g., if an SM can hold 32 blocks, shared mem per block 1280B, then per SM it's 40KB, well within 164KB. So occupancy is limited by warps: 1024 blocks * 2 warps = 2048 warps across 108 SMs => ~19 warps per SM. The A100 can hold 64 warps per SM, so occupancy is ~30%. That's okay. To improve occupancy, we could reduce shared memory usage, but it's already tiny. Actually occupancy is limited by the number of warps, not shared memory. If we increase block size to 128 threads (4 warps) per block, total warps = 4096, occupancy ~38%. But half the threads idle. Still, more warps. However, block size 128 with one node per block means many idle threads. The scheduler might still keep them, but they do no work. What if we process 2 nodes per block of 128 threads? Then each node gets 64 threads. Total blocks = 512. Total warps = 512 * 4 = 2048 warps (same as block size 64, one node per block). But now we have fewer blocks (512). Not a big difference. What if we process 4 nodes per block of 256 threads? Then each node gets 64 threads. Total blocks = 256. Total warps = 256 * 8 = 2048 warps. Even fewer blocks. As long as blocks > SM count, it should be fine. 256 blocks > 108 SMs, so each SM gets ~2-3 blocks. Since each block has 8 warps, each SM has 16-24 warps. Still okay. But fewer blocks might mean worse latency hiding if some SMs finish early. But the workload is uniform. I think one block per node with block size 64 is fine. Or maybe 128 to get more warps. Let's test with 128 threads per block, but each thread processes one node (grid = 1024/128 = 8 blocks). No, that's total threads = 1024, each thread one node. Warps=32. Too few. So we must have grid = B*N. So block size 64, grid = B*N => 1024 blocks, 2048 warps. Block size 128, grid = B*N => 1024 blocks, 4096 warps. More warps, but half idle. However, the warp scheduler simply skips inactive threads quickly. The main issue is that the active threads per warp might be half. For example, loading scalar_feat: 64 threads load, 64 threads don't. The warp with threads 64-127 will be idle. The warp with threads 0-63 will be active. So half the warps are idle. That means effective warps = 2048, same as block size 64. But we have more warps in total, but half do nothing. So occupancy is effectively the same. What about block size 32 (one warp per node)? Then all threads are active for scalar_in=64? No, loops with stride 32 and bound 64 mean each thread does 2 iterations. For vector_in*3=48, each thread does 1 or 2 iterations. For scalar_out=64, each thread does 2 iterations. For v_out=48, each thread does 1 or 2. All threads are active. Total warps = 1024. That's fewer warps than 2048. But all threads in each warp are active. Usually it's better to have more active warps per SM to hide latency. But 1024 warps across 108 SMs = ~9.5 warps per SM. Very low occupancy. That might hurt memory latency hiding. Block size 64: for scalar_in=64, half the threads (0..63) handle exactly one iteration, all active. For vector_in*3=48, all threads active (one iteration). For combined_size=80, copy scalar part: all active (one iteration). Norms: vector_in=16, threads 0..15 active, 16..63 idle? Actually our loop is `for (int i=tid; i<vector_in; i+=num_threads)`. With tid 0..63 and vector_in=16, only tid 0..15 enter the loop, 16..63 skip. So in a warp, threads 0..15 active, 16..31 idle (first warp), second warp (32..63) all idle. That's not great. Similarly for scalar_out=64, all threads active (one iteration). For vector_out=16, only threads 0..15 active. For v_out_size=48, threads 0..47 active, 48..63 idle. So there is some divergence within warps. Could we improve by remapping threads to tasks such that each thread does a subset of multiple tasks? Or by using a 2D thread layout? But for simplicity, block size 64 is okay. Alternatively, block size = 16? Then for scalar_in=64, each thread does 4 loads. For vector_in*3=48, each thread does 3 loads. For scalar_out=64, each thread does 4 accumulations. For vector_out=16, each thread does one gate. For v_out=48, each thread does 3 accumulations. All threads active. But total warps = 1024 blocks * 0.5 warps = 512 warps. That's very low. So block size 64 seems a decent tradeoff: for most tasks, a majority of threads are active. The only inactivity is for vector_in=16 and vector_out=16. That's one quarter of threads. Given the tiny workload, I think block size 64 is fine. One more optimization: we can store `combined` in registers rather than shared memory, because each thread that needs it (for scalar and gate) can just have a local copy. But since multiple threads need the same combined (each row of weight), shared memory is better because we load combined once from shared memory and broadcast to all threads. If we keep combined in registers, we'd need to load scalar_feat and norms into each thread's registers, which uses more registers. Shared memory is fine. Actually, can we avoid shared memory for combined and just have each thread compute the dot product directly from scalar_feat and vector_feat? Yes: - Each thread reads the scalar weights and bias for its output row. - It loads scalar_feat (64 floats) and vector_norms (16 floats) into local array (registers or local memory). Since combined is only 80 floats, we can keep it in local memory (spilled to local memory if registers insufficient). But with block size 64, and each thread needing 80 floats for combined, that's 5120 floats per block. If stored in registers, that's a lot (80 registers per thread). Might be okay. But local memory is slow. Using shared memory is much better. Let's refine the kernel to minimize syncthreads and shared memory. We need to store: - s_scalar: 64 floats - s_vector: 48 floats - s_combined: 80 floats - s_scalar_out: 64 floats - s_gate: 16 floats - s_vector_out: 48 floats Total 320 floats. We can allocate all in one shared memory block. Wait, do we need s_scalar_out and s_vector_out in shared memory? We could write directly to global memory. But since multiple threads write, and we do gating, we need s_vector_out or we can apply gate and write directly. Actually: - For scalar_out, each thread computes its own output and can write directly to global memory. No need for shared memory for scalar_out. - For gate, threads 0..15 compute gate values. They need to be shared with threads that apply gating to v_out. So we need s_gate (16 floats) in shared memory. - For v_out, threads compute vector linear results into shared memory (or registers?), then apply gate, then write to global memory. We can compute v_out, multiply by gate, and write to global memory all in one loop. But the gate is in shared memory. The vector linear result can be computed, multiplied by s_gate, and written directly without storing in shared memory. However, we have a __syncthreads() before that to ensure s_gate is ready. So we can avoid s_vector_out in shared memory. Similarly, s_scalar and s_vector are needed to compute combined. s_combined is needed for both scalar and gate. Could we avoid s_combined? Each thread for scalar and gate could read scalar_feat directly. But they need vector norms. Norms depend on s_vector. So we could compute norms and store them in shared memory (s_norm). Then scalar threads can load scalar_feat from shared memory and norms from shared memory. But s_scalar is already in shared memory. So s_combined is just a convenience. It's fine. Let's redesign to use minimal shared memory: ```cpp extern __shared__ char smem[]; scalar_t* s_scalar = (scalar_t*)smem; // scalar_in scalar_t* s_vector = s_scalar + scalar_in; // vector_in * 3 scalar_t* s_norm = s_vector + vector_in * 3; // vector_in scalar_t* s_gate = s_norm + vector_in; // vector_out ``` Total: 64+48+16+16 = 144 floats = 576 bytes. Even smaller. But then for scalar linear, each thread needs to compute dot(combined, w_row). Combined is scalar_in + vector_in = 80. It would need to read s_scalar[0..63] and s_norm[0..15]. That's two arrays. Fine: loop i=0..scalar_in-1: load s_scalar[i]; loop i=0..vector_in-1: load s_norm[i]. Similarly for gate. So we can eliminate s_combined and s_scalar_out and s_vector_out from shared memory. Let's write the kernel with this layout: 1. Load scalar_feat into s_scalar (tid loops). 2. Load vector_feat into s_vector (tid loops). 3. __syncthreads(). 4. Compute s_norm: for i = tid; i < vector_in; i += blockDim.x: x = s_vector[i*3], y = s_vector[i*3+1], z = s_vector[i*3+2]; s_norm[i] = sqrt(x*x + y*y + z*z); 5. __syncthreads(). 6. Scalar linear: for o = tid; o < scalar_out; o += blockDim.x: accum = scalar_bias[o]; // scalar part w_ptr = scalar_weight + o * combined_size; for i = 0..scalar_in-1: accum += w_ptr[i] * s_scalar[i]; for i = 0..vector_in-1: accum += w_ptr[scalar_in + i] * s_norm[i]; scalar_out_ptr[...] = accum; 7. Vector gate: for o = tid; o < vector_out; o += blockDim.x: accum = gate_bias[o]; w_ptr = gate_weight + o * combined_size; for i = 0..scalar_in-1: accum += w_ptr[i] * s_scalar[i]; for i = 0..vector_in-1: accum += w_ptr[scalar_in + i] * s_norm[i]; s_gate[o] = 1 / (1 + exp(-accum)); 8. __syncthreads(). 9. Vector linear + gate + write: int v_out_size = vector_out * 3; for idx = tid; idx < v_out_size; idx += blockDim.x: o = idx / 3; d = idx % 3; accum = 0; w_ptr = vector_weight + o * vector_in; for i = 0..vector_in-1: accum += w_ptr[i] * s_vector[i*3 + d]; accum *= s_gate[o]; vector_out_ptr[...] = accum; ``` This avoids s_scalar_out and s_vector_out. But we still need s_gate. Shared memory size is 144 floats = 576 bytes. One issue: step 6 writes scalar_out directly to global memory. That's fine. Step 7 writes s_gate. Step 9 reads s_gate and writes vector_out directly. Good. Now, we need to think about how `scalar_out` and `vector_out` writes are coalesced. For scalar_out, thread tid writes output channel tid. Adjacent threads write adjacent channels for the same node. In global memory, scalar_out layout is (B,N,scalar_out). For node_idx fixed, channel c is at offset `node_idx*scalar_out + c`. So adjacent threads (tid, tid+1) write adjacent memory. Coalesced. For vector_out, thread idx writes `idx = o*3 + d`. Similarly, adjacent threads write adjacent floats in the (vector_out*3) dimension. Coalesced. For vector_feat loads, thread tid reads `i = tid` for vector_feat. Layout of vector_feat is (B,N,vector_in,3). For a fixed node, the memory is arranged as vector_in groups of 3 floats: [v0_x, v0_y, v0_z, v1_x, v1_y, v1_z, ...]. So reading by linear index i across threads is contiguous. Perfect. For scalar_feat loads, similarly contiguous. For weights: - scalar_weight layout: (scalar_out, combined_size). Thread o reads row o. Address = o*combined_size + i. Adjacent threads (o, o+1) have stride combined_size. Not coalesced. But weights are read-only and likely cached. Since each row is read by one thread, it's contiguous per thread. Good for caches. Could store transposed weights to improve coalescing? Let's see: if we store `scalar_weight_T` of shape (combined_size, scalar_out), then thread o reads column o, which is strided by scalar_out. Still strided. Not helpful. We can keep as row-major. - vector_weight layout: (vector_out, vector_in). Thread idx reads row o. Similar. Given the tiny size, weight cache misses won't be a big issue. All weights fit in L2. One more optimization: use `float` for computations even if input is `half`. But inputs are fp32 (the problem says fp32). So we use float. We should also consider using `__ldg` for read-only data like weights and inputs, but modern compilers often do this automatically. We can mark pointers with `__restrict__` and `const` which we already did. Also, we should probably unroll inner loops for performance. Since sizes are dynamic but known at kernel launch, we can't unroll at compile time easily without templating. But we can use `#pragma unroll` for small loops with constant trip counts? Not if loop bound is variable. But the inner loops have bounds like `scalar_in` (64), `vector_in` (16), etc. We could template the kernel on these sizes. However, the sizes are passed via constructor. We could instantiate a generic kernel and a specialized one for common sizes, and dispatch at runtime. But for the benchmark, the sizes are fixed: 64,64,16,16. So we could just write a kernel that assumes these sizes? No, we must accept arbitrary constructor arguments. But we can create templated kernel with `scalar_in, scalar_out, vector_in, vector_out` as template parameters, and at runtime check if they match the benchmark sizes; if so, launch the specialized version; otherwise launch a generic version. The overhead of runtime check is trivial. The generic version can be written with loops. The specialized version can have `#pragma unroll` and fixed-size shared memory. Actually, we can write a generic kernel with `#pragma unroll` inside `#pragma unroll` - the compiler might unroll small loops if it can see the loop bound is a constant? No, it doesn't know. But we can still pass the sizes as `constexpr` if we template them: ```cpp template <int SIN, int SOUT, int VIN, int VOUT> __global__ void kernel(...) { ... } ``` Then in Python, we can write a dispatcher: ```python if scalar_in == 64 and scalar_out == 64 and vector_in == 16 and vector_out == 16: kernel = kernel_64_64_16_16 else: kernel = kernel_generic ``` Given the benchmark uses these exact sizes, the specialized kernel will run. Let's write the generic kernel first, then maybe a templated one. But writing both might be complicated. Let's just write the generic kernel and hope the compiler does decent loop optimization. The loops are small (up to 80). The overhead of loop control is negligible compared to memory latency. But for occupancy, register usage matters. A generic kernel with variable loops might use fewer registers? Actually loops with known small constant bounds allow compiler to unroll and use more registers. With variable bounds, compiler may keep them as loops. That could reduce register pressure and actually improve occupancy. But performance may be slightly lower. However, the generic kernel loops of size 80 and 16 are small; the compiler might fully unroll them if it sees the bounds are uniform across the kernel and small? It depends on the compiler's heuristics. Since the bounds are kernel parameters (passed as ints), it cannot unroll them. So it will generate loops. But loop execution overhead is tiny (80 iterations) inside each thread. The memory loads are the bottleneck. Wait, but the main cost in the generic kernel will be the loop overhead and the fact that the compiler can't optimize memory access patterns as aggressively. However, the memory loads are loads from shared memory, which are fast. The loads from global memory (weights) happen in the loops. Since weights are cached, the loop overhead is not a big deal. Let's try the generic kernel first and see if correctness passes and how fast it is. We also need to handle the `forward` method. The reference uses `nn.Linear` parameters: `scalar_linear.weight` and `.bias`, `vector_gate.weight` and `.bias`, plus `vector_weight`. In `ModelNew`, we can accept these constructor arguments and initialize parameters, or we can accept the module references. Actually the instructions say: "Accept the same constructor arguments as Model." So `ModelNew` takes `scalar_in, scalar_out, vector_in, vector_out`. It should probably define its own `nn.Parameter`s or `nn.Linear`s, or we can just store the weights and biases. But to be compatible with the reference, the `ModelNew` class should have the same parameter names? The evaluation code probably does something like: ```python model = ModelNew(*get_init_inputs()) inputs = get_inputs() out = model(*inputs) ``` And compare to `Model`. It may check outputs only. It may also initialize weights similarly. The problem says "Replace the forward computation of the PyTorch model below with a custom CUDA kernel." The constructor can set up parameters. To be safe, we should initialize parameters the same way as the reference. The reference does: ```python self.scalar_linear = nn.Linear(scalar_in + vector_in, scalar_out) self.vector_weight = nn.Parameter(torch.randn(vector_out, vector_in)) nn.init.xavier_uniform_(self.vector_weight) self.vector_gate = nn.Linear(scalar_in + vector_in, vector_out) ``` So in `ModelNew`, we should also define these parameters, but we can define them as `nn.Parameter` or as `nn.Linear` and then extract the weight/bias tensors to pass to the kernel. However, since we are replacing the forward pass, we can just define the same parameters so that the model struct matches. We can define: ```python class ModelNew(nn.Module): def __init__(self, scalar_in, scalar_out, vector_in, vector_out): super().__init__() self.scalar_linear = nn.Linear(scalar_in + vector_in, scalar_out) self.vector_weight = nn.Parameter(torch.randn(vector_out, vector_in)) nn.init.xavier_uniform_(self.vector_weight) self.vector_gate = nn.Linear(scalar_in + vector_in, vector_out) ``` Then in `forward`, we call the custom CUDA extension with `self.scalar_linear.weight`, `self.scalar_linear.bias`, `self.vector_weight`, `self.vector_gate.weight`, `self.vector_gate.bias`. We need to compile the extension at module load time using `torch.utils.cpp_extension.load_inline`. The C++ code shouldn't include a pybind module. Let's write the inline C++ code. C++ source: ```cpp #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> template <typename scalar_t> __global__ void se3_forward_kernel( const scalar_t* __restrict__ scalar_feat, const scalar_t* __restrict__ vector_feat, const scalar_t* __restrict__ scalar_weight, const scalar_t* __restrict__ scalar_bias, const scalar_t* __restrict__ vector_weight, const scalar_t* __restrict__ gate_weight, const scalar_t* __restrict__ gate_bias, scalar_t* __restrict__ scalar_out, scalar_t* __restrict__ vector_out, int B, int N, int scalar_in, int scalar_out, int vector_in, int vector_out) { int node_idx = blockIdx.x; int tid = threadIdx.x; int nt = blockDim.x; if (node_idx >= B * N) return; int combined_size = scalar_in + vector_in; int vfeat_size = vector_in * 3; int vout_size = vector_out * 3; extern __shared__ char smem[]; scalar_t* s_scalar = reinterpret_cast<scalar_t*>(smem); scalar_t* s_vector = s_scalar + scalar_in; scalar_t* s_norm = s_vector + vfeat_size; scalar_t* s_gate = s_norm + vector_in; // load scalar_feat for (int i = tid; i < scalar_in; i += nt) { s_scalar[i] = scalar_feat[node_idx * scalar_in + i]; } // load vector_feat for (int i = tid; i < vfeat_size; i += nt) { s_vector[i] = vector_feat[node_idx * vfeat_size + i]; } __syncthreads(); // compute norms for (int i = tid; i < vector_in; i += nt) { scalar_t x = s_vector[i * 3]; scalar_t y = s_vector[i * 3 + 1]; scalar_t z = s_vector[i * 3 + 2]; s_norm[i] = sqrtf(x*x + y*y + z*z); } __syncthreads(); // scalar linear const scalar_t* w_s = scalar_weight; for (int o = tid; o < scalar_out; o += nt) { scalar_t acc = scalar_bias[o]; const scalar_t* row = w_s + o * combined_size; #pragma unroll for (int i = 0; i < scalar_in; ++i) { acc += row[i] * s_scalar[i]; } #pragma unroll for (int i = 0; i < vector_in; ++i) { acc += row[scalar_in + i] * s_norm[i]; } scalar_out[node_idx * scalar_out + o] = acc; } // vector gate const scalar_t* w_g = gate_weight; for (int o = tid; o < vector_out; o += nt) { scalar_t acc = gate_bias[o]; const scalar_t* row = w_g + o * combined_size; #pragma unroll for (int i = 0; i < scalar_in; ++i) { acc += row[i] * s_scalar[i]; } #pragma unroll for (int i = 0; i < vector_in; ++i) { acc += row[scalar_in + i] * s_norm[i]; } scalar_t g = 1.0f / (1.0f + expf(-acc)); s_gate[o] = g; } __syncthreads(); // vector linear + gate const scalar_t* w_v = vector_weight; for (int idx = tid; idx < vout_size; idx += nt) { int o = idx / 3; int d = idx % 3; scalar_t acc = 0.0f; const scalar_t* row = w_v + o * vector_in; #pragma unroll for (int i = 0; i < vector_in; ++i) { acc += row[i] * s_vector[i * 3 + d]; } acc *= s_gate[o]; vector_out[node_idx * vout_size + idx] = acc; } } ``` Wait, `#pragma unroll` with variable loop bounds is ignored by the compiler, but it doesn't cause errors. We can leave it. Actually, we can't easily unroll without template parameters. But maybe we can add a specialized path. However, note: `scalar_in + vector_in` may not be known at compile time. `#pragma unroll` won't work. We can still use it but it does nothing. One concern: using `1.0f / (1.0f + expf(-acc))` is fine for float. For double, use `exp`. Since inputs are fp32, `float` is fine. But `scalar_t` might be `float` and `scalar_t(1.0f)` is okay. The expression `1.0f / (1.0f + expf(-acc))` works for `float` but if `scalar_t` is `double`, `expf` takes float. We should use `exp` or `expf` conditionally. Since inputs are fp32, we can just assume float. But to be generic, we can write: ```cpp scalar_t g = scalar_t(1) / (scalar_t(1) + exp(-acc)); ``` Wait, `exp` is overloaded for float and double in C++? `exp` from `<cmath>` has overloads. But on CUDA, `exp` for `float` is not defined? Actually `exp(float)` is available if we include `<cmath>` or use `::exp`. But in device code, `expf` is the standard CUDA intrinsic for float, and `exp` for double. To be safe, we can write: ```cpp scalar_t g = scalar_t(1) / (scalar_t(1) + static_cast<scalar_t>(1) / (static_cast<scalar_t>(1) + exp(-acc))); ``` Actually we can define an inline helper: ```cpp template <typename T> __device__ __forceinline__ T sigmoid(T x) { return T(1) / (T(1) + exp(-x)); } ``` But `exp` for float in device code might resolve to `::exp(float)` which is supported. Or we can specialize: ```cpp __device__ __forceinline__ float sigmoid(float x) { return 1.0f / (1.0f + expf(-x)); } __device__ __forceinline__ double sigmoid(double x) { return 1.0 / (1.0 + exp(-x)); } ``` That is cleaner. Now, the C++ host wrapper: ```cpp std::vector<torch::Tensor> se3_forward( torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_weight, torch::Tensor scalar_bias, torch::Tensor vector_weight, torch::Tensor gate_weight, torch::Tensor gate_bias) { int B = scalar_feat.size(0); int N = scalar_feat.size(1); int scalar_in = scalar_feat.size(2); int vector_in = vector_feat.size(2); int scalar_out = scalar_weight.size(0); int vector_out = vector_weight.size(0); auto scalar_out_t = torch::empty({B, N, scalar_out}, scalar_feat.options()); auto vector_out_t = torch::empty({B, N, vector_out, 3}, vector_feat.options()); int nodes = B * N; int threads = 64; int combined_size = scalar_in + vector_in; int vfeat_size = vector_in * 3; int shared_mem_size = (scalar_in + vfeat_size + vector_in + vector_out) * sizeof(float); AT_DISPATCH_FLOATING_TYPES(scalar_feat.scalar_type(), "se3_forward", [&] { se3_forward_kernel<scalar_t><<<nodes, threads, shared_mem_size>>>( scalar_feat.data_ptr<scalar_t>(), vector_feat.data_ptr<scalar_t>(), scalar_weight.data_ptr<scalar_t>(), scalar_bias.data_ptr<scalar_t>(), vector_weight.data_ptr<scalar_t>(), gate_weight.data_ptr<scalar_t>(), gate_bias.data_ptr<scalar_t>(), scalar_out_t.data_ptr<scalar_t>(), vector_out_t.data_ptr<scalar_t>(), B, N, scalar_in, scalar_out, vector_in, vector_out); }); return {scalar_out_t, vector_out_t}; } ``` But wait: `scalar_feat.options()` might be CPU/CUDA? We need to ensure it's on CUDA. The inputs are likely on GPU because the benchmark runs on GPU. `load_inline` compiles CUDA code. We need to make sure the tensors are contiguous. The `norm` in reference computes `sqrt(sum(x^2))`. Our kernel uses `sqrtf(x*x + y*y + z*z)`. That's correct. One subtlety: `torch.cat` produces a contiguous tensor? Yes. The linear layers use the concatenated dimension. Our kernel replicates that. Let's also consider templated specialization for the exact sizes to allow loop unrolling. Since the sizes are small, unrolling might improve performance slightly, but the compiler with `#pragma unroll` on variable loops may still not unroll. However, the loop overhead is minimal (128 iterations max per dot product). The main cost is memory loads. But we can write a specialized version: ```cpp template <typename scalar_t> __global__ void se3_forward_kernel_specialized_64_64_16_16(...) { // same code but with #pragma unroll(64) etc. or manual unroll? } ``` Actually we can just write the generic version and then a specialized version using loops with constants. But the generic version should be fine. Let's test with generic. One more thing: `extern __shared__ char smem[]` requires dynamic shared memory. The size is passed at kernel launch. Since the kernel is generic, we compute `shared_mem_size` based on the input dimensions. But `reinterpret_cast<scalar_t*>(smem)` may need alignment. `char` arrays are aligned to 1. `scalar_t` (float) requires alignment 4. But dynamic shared memory is aligned to the largest type? Actually CUDA ensures dynamic shared memory is aligned to 16 bytes or to the alignment of the type used in the first cast. But to be safe, we can use `__shared__ float smem[]` directly, but we need the type to be scalar_t. We can do `extern __shared__ float smem_f[];` and reinterpret. Or we can use `extern __shared__ __align__(sizeof(float)) char smem[];`. Actually, C++ doesn't allow alignment attributes on `extern` arrays? We can just allocate as `extern __shared__ float smem[];` and then cast to scalar_t*. Since scalar_t is float/double, `float` is okay. For double, `extern __shared__ double smem[];` would be needed if scalar_t is double. But we can just use `extern __shared__ char smem[];` and the compiler aligns it to the natural alignment of the requested type? Actually, dynamic shared memory base is aligned to 16 bytes on modern GPUs. So `reinterpret_cast<scalar_t*>(smem)` is safe. Let's use `extern __shared__ float smem[];` for simplicity, since we know it's float. But to be generic for double, we can do `extern __shared__ char smem[];`. Actually, `AT_DISPATCH_FLOATING_TYPES` covers float and double. The benchmark uses float. Let's just use `char` and rely on 16-byte alignment provided by CUDA runtime for dynamic shared memory. Wait, is dynamic shared memory base aligned to 16 bytes? Yes, on all compute capabilities >= 2.x. So cast to float* is safe. Now, about `vector_feat` shape: (B, N, vector_in, 3). We treat it as flattened `(B*N, vector_in*3)`. The memory layout is contiguous in that dimension. One potential issue: the reference uses `nn.Linear` which includes bias. Our kernel uses `scalar_bias`. Good. Another potential issue: the reference initializes `self.vector_weight` with `xavier_uniform_`. That's fine. We just use the parameter. Now, we need `ModelNew` class in Python. ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline # define cpp source and cuda source cpp_source = """ torch::Tensor se3_forward(...); """ # Actually, we need to define the function and then load it. ``` Let's write the full `kernel_code` string. We need to be careful about quoting and indentation. The `load_inline` call should be at module level (outside class). We define a function `se3_forward` that calls the kernel. Let's draft the code. ```python import torch import torch.nn as nn cuda_source = r""" #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> template <typename T> __device__ __forceinline__ T sigmoid_device(T x) { return T(1) / (T(1) + exp(-x)); } template <> __device__ __forceinline__ float sigmoid_device(float x) { return 1.0f / (1.0f + expf(-x)); } template <> __device__ __forceinline__ double sigmoid_device(double x) { return 1.0 / (1.0 + exp(-x)); } template <typename scalar_t> __global__ void se3_forward_kernel( const scalar_t* __restrict__ scalar_feat, const scalar_t* __restrict__ vector_feat, const scalar_t* __restrict__ scalar_weight, const scalar_t* __restrict__ scalar_bias, const scalar_t* __restrict__ vector_weight, const scalar_t* __restrict__ gate_weight, const scalar_t* __restrict__ gate_bias, scalar_t* __restrict__ scalar_out, scalar_t* __restrict__ vector_out, int B, int N, int scalar_in, int scalar_out, int vector_in, int vector_out) { int node_idx = blockIdx.x; int tid = threadIdx.x; int nt = blockDim.x; if (node_idx >= B * N) return; int combined_size = scalar_in + vector_in; int vfeat_size = vector_in * 3; int vout_size = vector_out * 3; extern __shared__ char smem[]; scalar_t* s_scalar = reinterpret_cast<scalar_t*>(smem); scalar_t* s_vector = s_scalar + scalar_in; scalar_t* s_norm = s_vector + vfeat_size; scalar_t* s_gate = s_norm + vector_in; // Load scalar features for (int i = tid; i < scalar_in; i += nt) { s_scalar[i] = scalar_feat[node_idx * scalar_in + i]; } // Load vector features for (int i = tid; i < vfeat_size; i += nt) { s_vector[i] = vector_feat[node_idx * vfeat_size + i]; } __syncthreads(); // Compute vector norms for (int i = tid; i < vector_in; i += nt) { scalar_t x = s_vector[i * 3]; scalar_t y = s_vector[i * 3 + 1]; scalar_t z = s_vector[i * 3 + 2]; s_norm[i] = sqrtf(x*x + y*y + z*z); } __syncthreads(); // Scalar linear const scalar_t* w_s = scalar_weight; for (int o = tid; o < scalar_out; o += nt) { scalar_t acc = scalar_bias[o]; const scalar_t* row = w_s + o * combined_size; for (int i = 0; i < scalar_in; ++i) { acc += row[i] * s_scalar[i]; } for (int i = 0; i < vector_in; ++i) { acc += row[scalar_in + i] * s_norm[i]; } scalar_out[node_idx * scalar_out + o] = acc; } // Vector gate const scalar_t* w_g = gate_weight; for (int o = tid; o < vector_out; o += nt) { scalar_t acc = gate_bias[o]; const scalar_t* row = w_g + o * combined_size; for (int i = 0; i < scalar_in; ++i) { acc += row[i] * s_scalar[i]; } for (int i = 0; i < vector_in; ++i) { acc += row[scalar_in + i] * s_norm[i]; } s_gate[o] = sigmoid_device(acc); } __syncthreads(); // Vector linear + gating const scalar_t* w_v = vector_weight; for (int idx = tid; idx < vout_size; idx += nt) { int o = idx / 3; int d = idx % 3; scalar_t acc = 0.0f; const scalar_t* row = w_v + o * vector_in; for (int i = 0; i < vector_in; ++i) { acc += row[i] * s_vector[i * 3 + d]; } acc *= s_gate[o]; vector_out[node_idx * vout_size + idx] = acc; } } torch::Tensor se3_forward( torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_weight, torch::Tensor scalar_bias, torch::Tensor vector_weight, torch::Tensor gate_weight, torch::Tensor gate_bias) { int B = scalar_feat.size(0); int N = scalar_feat.size(1); int scalar_in = scalar_feat.size(2); int vector_in = vector_feat.size(2); int scalar_out = scalar_weight.size(0); int vector_out = vector_weight.size(0); auto scalar_out_t = torch::empty({B, N, scalar_out}, scalar_feat.options()); auto vector_out_t = torch::empty({B, N, vector_out, 3}, vector_feat.options()); int nodes = B * N; int threads = 64; int vfeat_size = vector_in * 3; int shared_mem_size = (scalar_in + vfeat_size + vector_in + vector_out) * sizeof(float); AT_DISPATCH_FLOATING_TYPES(scalar_feat.scalar_type(), "se3_forward", [&] { se3_forward_kernel<scalar_t><<<nodes, threads, shared_mem_size>>>( scalar_feat.data_ptr<scalar_t>(), vector_feat.data_ptr<scalar_t>(), scalar_weight.data_ptr<scalar_t>(), scalar_bias.data_ptr<scalar_t>(), vector_weight.data_ptr<scalar_t>(), gate_weight.data_ptr<scalar_t>(), gate_bias.data_ptr<scalar_t>(), scalar_out_t.data_ptr<scalar_t>(), vector_out_t.data_ptr<scalar_t>(), B, N, scalar_in, scalar_out, vector_in, vector_out); }); return {scalar_out_t, vector_out_t}; } """ ``` Actually, the C++ wrapper returns `std::vector<torch::Tensor>`, but we need to match the pybind return type. In ATen, returning a vector of tensors from a function registered with `m.def` works. But in `load_inline`, we list the function name in `functions=[...]`. The function signature should be `std::vector<torch::Tensor> se3_forward(...)`. That's fine. One issue: in the CUDA source, `#include <math.h>` might conflict. Better to use `<cmath>` or just rely on CUDA builtins. `sqrtf`, `expf` are available. Another issue: `torch::Tensor se3_forward(...)` returns a vector but the return type is `torch::Tensor`. Wait, in the snippet above I wrote `torch::Tensor se3_forward` but the function returns `{scalar_out_t, vector_out_t}`. In C++, `torch::Tensor` is a single tensor, so returning an initializer list with two tensors will cause a compile error. We need `std::vector<torch::Tensor> se3_forward(...)`. Let's fix that. Also, in the kernel, `sqrtf` is for float. If scalar_t is double, we should use `sqrt`. To be generic, we can write `sqrt(x*x + y*y + z*z)` and rely on overload resolution. CUDA provides `sqrt(double)` and `sqrtf(float)`. For `scalar_t`, we need to call the right one. Since scalar_t is either float or double, we can use `sqrt` which is overloaded for both in `<cmath>`. But in device code, `sqrt` is a built-in for double, and for float it's `sqrtf` or `sqrt(float)`. Actually CUDA supports overloaded `sqrt` for float and double? I think yes: `sqrt(float)` maps to `sqrtf`, `sqrt(double)` maps to `sqrt`. To be safe, we can write: ```cpp s_norm[i] = sqrt(x*x + y*y + z*z); ``` If `x` is `float`, `sqrt(float)` calls `sqrtf`. If `x` is `double`, it calls `sqrt(double)`. This is fine. Similarly, `exp(-acc)` inside `sigmoid_device`: we need overloaded `exp`. For `float`, `exp(float)` calls `expf`. For `double`, `exp(double)`. So we can write `exp(-x)`. Let's implement it as: ```cpp template <typename T> __device__ __forceinline__ T sigmoid_device(T x) { return T(1) / (T(1) + exp(-x)); } ``` But `exp` is overloaded in CUDA's `::` namespace. In device code, `exp` is a built-in function overloaded for float, double, half, etc. So `exp(-x)` should work. However, sometimes there are issues with `exp` from `<cmath>` being host-only. To avoid that, we can use `std::exp` but that's not available in device code. The CUDA built-in `exp` is in the global namespace. So just `exp(-x)` is fine. But there is a potential ambiguity if `<math.h>` is included. Let's avoid `<math.h>` and just let CUDA built-ins handle it. Also, for `sqrt`, we can just write `sqrt(x*x + y*y + z*z)`. The CUDA built-in `sqrt` is overloaded. Let's also set `blockDim.x` to something like 64. But we can tune this. Maybe 32 is enough? Let's think: scalar_in=64, scalar_out=64, vector_in=16, vector_out=16, vout_size=48. With 64 threads: - scalar loads: 64 threads, each loads 1 -> 1 transaction per warp? Actually 32 threads load 32 adjacent floats = 128 bytes, perfect coalescing. 2 warps, 2 transactions. - vector loads: 48 floats. First warp loads 32 floats (128 bytes), second warp loads 16 floats (64 bytes). Coalesced. - norms: 16 threads active. First 16 threads of warp 0 active. Loads from shared memory, no coalescing issue. - scalar linear: 64 threads active. Each thread reads weight row (80 floats). Weight reads are not coalesced across threads, but cached. - gate: 16 threads active. - vector linear: 48 threads active (threads 0..47). Each reads weight row (16 floats). With 32 threads: - scalar loads: 32 threads, each loads 2 -> coalesced. - vector loads: 32 threads, each loads 1 or 2 -> coalesced. - norms: 16 threads active in warp 0. - scalar linear: 32 threads, each does 2 outputs. - gate: 16 threads active. - vector linear: 32 threads, each does 1 or 2 outputs. 32 threads might be better because all active threads are in one warp, reducing warp count but fully utilizing that warp. However, with 64 threads, some operations have idle threads in warp 1. The warp scheduler still schedules warp 1 even if some threads are predicated off. But instructions are still issued. It wastes some instruction issue bandwidth but not a huge deal. However, block size 32 gives 1 warp per block. 1024 blocks => 1024 warps. Block size 64 gives 2 warps per block => 2048 warps. More warps generally help hide latency. But since the kernel is very short, latency hiding might not matter. I think block size 64 is a good start. But we can also try block size 128: for scalar loads, 64 threads active, 64 idle. For vector loads, 48 active, 80 idle. For scalar linear, 64 active, 64 idle. For vector linear, 48 active, 80 idle. So effective warps are same as 64 threads? Not exactly: a warp is either partially active or fully active. Warp 0 (threads 0..31) might have all active for scalar linear? Actually scalar linear loops over o = tid; o < scalar_out; o += nt. With nt=128 and scalar_out=64, only tid 0..63 execute. So warp 0 (0..31) and warp 1 (32..63) are fully active. Warps 2 and 3 (64..127) are fully inactive. So the scheduler will quickly skip warps 2 and 3. So effective instruction issue is same as block size 64. But total warps = 4096 vs 2048. The scheduler may waste time scanning inactive warps. So block size 64 is better. Thus, I will stick with 64 threads. One thing: `threads = 64` is hardcoded in the host wrapper. For generic sizes, we might want to choose threads based on e.g. `max(scalar_out, vector_out*3, ...)`. But we can just fix 64 for simplicity; it works for the given sizes. If scalar_out is larger than 64, the loops will handle it with stride. If scalar_in is huge, the shared memory might be large. But the benchmark uses small sizes. Let's add a check in the host wrapper to ensure block size is enough for coalescing? Not needed. Let's refine the code: ```cpp std::vector<torch::Tensor> se3_forward_cuda( torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_weight, torch::Tensor scalar_bias, torch::Tensor vector_weight, torch::Tensor gate_weight, torch::Tensor gate_bias) { int B = scalar_feat.size(0); int N = scalar_feat.size(1); int scalar_in = scalar_feat.size(2); int vector_in = vector_feat.size(2); int scalar_out = scalar_weight.size(0); int vector_out = vector_weight.size(0); auto scalar_out_t = torch::empty({B, N, scalar_out}, scalar_feat.options()); auto vector_out_t = torch::empty({B, N, vector_out, 3}, vector_feat.options()); int nodes = B * N; int threads = 64; int vfeat_size = vector_in * 3; int shared_mem_size = (scalar_in + vfeat_size + vector_in + vector_out) * sizeof(float); AT_DISPATCH_FLOATING_TYPES(scalar_feat.scalar_type(), "se3_forward_cuda", [&] { se3_forward_kernel<scalar_t><<<nodes, threads, shared_mem_size>>>( scalar_feat.data_ptr<scalar_t>(), vector_feat.data_ptr<scalar_t>(), scalar_weight.data_ptr<scalar_t>(), scalar_bias.data_ptr<scalar_t>(), vector_weight.data_ptr<scalar_t>(), gate_weight.data_ptr<scalar_t>(), gate_bias.data_ptr<scalar_t>(), scalar_out_t.data_ptr<scalar_t>(), vector_out_t.data_ptr<scalar_t>(), B, N, scalar_in, scalar_out, vector_in, vector_out); }); return {scalar_out_t, vector_out_t}; } ``` In Python: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline cpp_source = "" # we only need the host function in cuda_source, but load_inline expects cpp_sources for host and cuda_sources for device. Actually we can put everything in cuda_sources and have cpp_sources empty or with declarations. # Actually, we can put the host wrapper and kernel in cuda_sources, and set cpp_sources to "torch::Tensor se3_forward_cuda(...);" or just "" and use functions argument. ``` Wait, `load_inline` has arguments: - `cpp_sources`: C++ source code (host side) - `cuda_sources`: CUDA source code (device side) - `functions`: list of function names to expose We can put the entire code in `cuda_sources`. The host function is also compiled by nvcc, so it's fine. `cpp_sources` can be empty. Actually, `torch.utils.cpp_extension.load_inline` expects both. If we put everything in `cuda_sources`, it's okay. `cpp_sources` can be empty string or contain `#include <torch/extension.h>`. But there is a subtlety: when the host function returns a `std::vector<torch::Tensor>`, the pybind11 wrapper generated by `load_inline` must know how to handle it. `load_inline` auto-generates wrappers for functions listed in `functions`. It uses `pybind11` to bind them. For `std::vector<torch::Tensor>`, pybind11 should automatically convert to a Python tuple/list of torch tensors? Let's check. In PyTorch's C++ extension examples, returning `std::vector<torch::Tensor>` is common and works. The pybind11 wrapper returns a `py::list` or `std::vector<torch::Tensor>` which maps to a Python list. In the Python code, we can unpack it. Let's look at standard usage: ```python my_op = load_inline(name="my_op", cuda_sources=[...], functions=["my_func"]) out1, out2 = my_func(...) ``` Actually, `load_inline` returns a module. `module.my_func()` returns the outputs. If the C++ function returns `std::vector<torch::Tensor>`, pybind11 maps it to a Python list. So we can do: ```python tmp = _ext.se3_forward_cuda(...) return tmp[0], tmp[1] ``` Yes. Potential issue: If we put everything in `cuda_sources`, the host code is compiled by nvcc as a .cu file, which can handle C++ code. That's fine. Let's write the complete kernel_code. One more thought: We should ensure that the tensors are contiguous. In `forward`, we can call `.contiguous()` on inputs. But the reference likely receives contiguous tensors from `get_inputs()`. To be safe, we can call `.contiguous()`. Also, `scalar_weight` from `nn.Linear` is stored as `(out_features, in_features)`, which is the same as we expect: row-major, each row is weight for one output. `nn.Linear.weight` is actually of shape `(out_features, in_features)`. So `scalar_weight.size(0) = scalar_out`, `scalar_weight.size(1) = scalar_in + vector_in`. Good. `vector_weight` is shape `(vector_out, vector_in)`. Good. In the kernel, we assume row-major. PyTorch tensors are row-major (except for strides). If they are contiguous, `.data_ptr()` gives the correct layout. Let's also handle the case where inputs might be on CPU. We can move them to CUDA? The reference model likely runs on GPU if we create it with `.cuda()`. But the evaluation code may instantiate `ModelNew` and call it with inputs from `get_inputs()`. It may not call `.cuda()`. However, the performance comparison expects GPU execution. Usually these tasks assume GPU. The reference `Model` is a plain nn.Module; the evaluation likely moves it and inputs to the GPU. Since our `ModelNew` will be moved to GPU by `model.to(device)`, the parameters will be CUDA tensors. The inputs will be CUDA tensors. So our kernel will receive CUDA tensors. If inputs are CPU, the kernel launch will fail. But we can add a check in Python: ```python if not scalar_feat.is_cuda: raise RuntimeError("Inputs must be CUDA tensors") ``` But not necessary. Now, let's check correctness numerically. Potential sources of difference: 1. `norm` computation: `vector_feat.norm(dim=-1)` uses `sqrt(sum(x**2))` with default `p=2`. Our kernel does the same. However, `torch.norm` may use a numerically stable algorithm or `hypot`? Actually `torch.norm` for vectors is `sqrt(sum(x**2))`. It might be implemented as `sqrt` of `sum` of squares. The same. 2. Sigmoid: `torch.sigmoid` uses `1/(1+exp(-x))` but might use a stable formulation for large x, e.g., `exp(x)/(1+exp(x))` for x>0. For correctness tolerance 1e-4, our naive sigmoid should be fine. But let's check: if x is large negative, `exp(-x)` overflows to inf, so `1/(1+inf)` = 0. That's fine. If x is large positive, `exp(-x)` underflows to 0, so `1/(1+0)` = 1. That's fine. For extreme inputs, `torch.sigmoid` might not overflow because it uses different branches. But random inputs from `get_inputs()` have mean 0, std 1. After linear transforms, the gate pre-activation could be large? The weights are initialized with xavier and biases are default small (uniform). The input is standard normal. The output of a linear layer with ~80 inputs has variance ~1, so pre-activation is ~N(0,1). So `exp(-x)` won't overflow/underflow for typical values. So naive sigmoid is fine. 3. Order of operations: We compute vector linear before applying gate. The reference does `v_out = einsum(...)` then `gate = sigmoid(...)` then `v_out = gate * v_out`. Our kernel does the same. But we compute gate in parallel with scalar_out, then apply it to v_out. Good. 4. Float precision: We accumulate in `scalar_t`. The reference uses fp32. Our kernel uses fp32. The dot product accumulation order may differ because PyTorch's linear uses cuBLAS which does tree reductions in registers/shared memory. Our kernel accumulates serially per thread. For 80 terms, the difference should be small, well within 1e-4. Let's verify: typical dot product of 80 terms, each ~1, sum ~80. Serial accumulation vs split accumulation may differ by ~1e-6 relative. So tolerance 1e-4 is fine. 5. We must make sure the shapes are correct. Edge Cases: - `combined_size` dynamic. Works. - `B*N` large. Works. Let's compile and test. Actually, before writing the final code, let's create a draft and compile/correctness check.

Let me analyze the reference model: 1. `vector_feat` has shape `(B, N, vector_in, 3)` 2. `vec_norm = vector_feat.norm(dim=-1)` gives `(B, N, vector_in)` 3. `combined = torch.cat([scalar_feat, vec_norm], dim=-1)` gives `(B, N, scalar_in + vector_in)` 4. `s_out = self.scalar_linear(combined)` - linear layer 5. `v_out = torch.einsum("oi,bnid->bnod", self.vector_weight, vector_feat)` - this is a batched matrix multiply 6. `gate = torch.sigmoid(self.vector_gate(combined)).unsqueeze(-1)` 7. `v_out = gate * v_out` Let me compute the dimensions: - B = 8, N = 128 - scalar_in = 64, scalar_out = 64 - vector_in = 16, vector_out = 16 For the scalar linear: - Input: 64 + 16 = 80, Output: 64 - Weights: 80 x 64, bias: 64 For the vector gate: - Input: 80, Output: 16 - Weights: 80 x 16, bias: 16 For vector_weight: - 16 x 16 The einsum "oi,bnid->bnod" means for each b, n, multiply vector_weight (16x16) with vector_feat (16x3) to get (16x3). This is essentially a matrix multiplication per (b,n). The key operations are: 1. Compute L2 norm of vectors: sqrt(sum(x^2)) 2. Concatenate scalars with norms 3. Linear layer for s_out 4. Linear layer for gate 5. Matrix multiply vector_weight with vector_feat 6. Sigmoid activation 7. Multiply gate with v_out The main optimization opportunities: 1. Fuse the norm computation with concatenation and the linear layers 2. Fuse the vector transformation with gating 3. Use custom CUDA kernel to avoid intermediate allocations 4. Take advantage of the small sizes to do everything in registers/shared memory Let me think about kernel design. The sizes are: - B*N = 8*128 = 1024 independent operations - For scalar path: each "row" processes 80 inputs to 64+16=80 outputs - For vector path: each "row" does 16x16 MM with 16x3 Since each (b,n) is independent, we can parallelize over B*N. For each (b,n): 1. Compute norms of 16 vectors of dimension 3 2. Concatenate with 64 scalars 3. Compute scalar linear (80 -> 64) 4. Compute gate linear (80 -> 16) + sigmoid 5. Compute v_out = W_v @ v_feat (16x16 @ 16x3 = 16x3) 6. Scale v_out by gate This can all be done in one kernel! We load state, compute, and store results. Let me think about thread assignment: - We have B*N = 1024 items - Each item needs to compute: - Norms (16 * 3 = 48 floats) - Scalar linear: effectively a mat-vec of 80x64 - Gate linear: mat-vec of 80x16 - Vector transform: 16x16 @ 16x3 - Elementwise sigmoid and multiply We could assign one warp (32 threads) per (b,n) item, or multiple items per block. Let me think more carefully: - A block of 128 threads could handle 4 (b,n) items with 32 threads each - Or a block of 256 threads could handle 8 items For the linear layers with small dimensions: - scalar_linear: 80x64 matrix - vector_gate: 80x16 matrix - vector_weight: 16x16 matrix Loading weights from shared memory would be ideal. We can load all weights into shared memory once per block. Weights sizes: - scalar_weight: 80*64 = 5120 floats ≈ 20KB - scalar_bias: 64 floats - gate_weight: 80*16 = 1280 floats ≈ 5KB - gate_bias: 16 floats - vector_weight: 16*16 = 256 floats ≈ 1KB - Total weights: ~26KB - fits in shared memory easily! But B*N = 1024, and weights are shared across all. If we load weights in shared mem, we need to coordinate. Actually, let me think about this differently. Using shared memory for weights is great if multiple threads in a block use them. Option 1: Each block has multiple (b,n) items processed by warps, weights in shared memory. Option 2: Each thread handles one (b,n) item, registers for everything. Option 2 might work well because: - The input per item is 64 + 16*3 = 112 floats - The computation is: 2 linear layers + 1 matmul + elementwise + norms - For H100, we have plenty of registers Actually, let me think about registers. Per thread for one (b,n): - Input scalars: 64 floats → 256 bytes in registers is too much - Actually we don't need all in registers at once. We can stream through. For the scalar_linear output: 64 floats. We could compute these one or a few at a time. Better approach: use 2D tiling or have each warp handle one (b,n). Let me design a kernel where each thread block handles multiple (b,n), and within each block, threads cooperate. For scalar_linear with out_features=64 and in_features=80: - We can have 64 threads each compute one output element using a dot product of the input with one column of weight. Actually, more standard approach: - For each (b,n), we need to compute: y[j] = sum_i(x[i] * W[i,j]) + b[j] - We can assign T threads per (b,n), where each thread computes a slice of outputs. With 80 inputs and 64 outputs for scalar, and 16 outputs for gate: - If we assign 32 threads per (b,n): - Thread t computes output indices t, t+32, t+64, ... for scalar (up to 64) - And same for gate (up to 16) Wait, 64 outputs with 32 threads: each thread does 2 outputs for scalar linear. Gate has 16 outputs: each thread does 1 output (some threads do 0 or we gang them). Hmm, but for vector transformation: 16x16 @ 16x3 = 16x3. We need to compute 16 output vectors, each of dim 3 = 48 floats. With 32 threads: each thread does 1.5 outputs... doesn't divide evenly. 16x3 = 48 elements. With 32 threads: 16 threads do 2 elements, 16 threads do 1? No, for matrix multiply each thread can compute one element of output. 48 outputs / 32 threads = 1.5 per thread. Not clean. What if we use 16 threads per (b,n)? - Scalar: 64 outputs / 16 = 4 outputs per thread - Gate: 16 outputs / 16 = 1 output per thread - Vector MM: 48 outputs / 16 = 3 outputs per thread That divides evenly! But with 16 threads, we have less parallelism within each (b,n). Since B*N = 1024, total threads = 1024 * 16 = 16384. H100 can handle this. What about 48 threads per (b,n)? - Scalar: 64/48 - doesn't divide What about 64 threads? - Scalar: 1 per thread - Gate: 16 threads do 1, 48 do nothing (or we reduce to 16 active) - Vector: 48/64 - not clean Actually, we can use warp-level primitives. A warp of 32 threads: - For 64-output linear: each thread computes 2 outputs (reduction across warp for dot product) - For 16-output linear: each thread computes 1 output (with 16 threads active, or all 32 with 2-way reduction) - For 16x16x3 = 48 elements: each thread computes 1 or 2 elements, with 16 threads computing 2 and 16 computing 1? Hmm. Actually, for the matrix multiply C = A @ B where A is 16x16 and B is 16x3: C[i,j] = sum_k(A[i,k] * B[k,j]) With 32 threads in warp: Option: each thread computes some (i,j) pairs. 48 elements / 32 threads = 1.5 - not clean. With 16 threads: 48/16 = 3 clean. Each thread computes 3 elements. Let me think about using 16 threads per (b,n): Total parallel work: 1024 * 16 = 16384 threads. Block size: 128 threads = 8 (b,n) per block. Number of blocks: 1024/8 = 128 blocks. In each block, 8 (b,n) items, each handled by 16 threads. For linear layers, we need to load inputs and weights. Let me think about using shared memory for weights: - scalar_weight: 80 x 64 = 5120 floats = 20KB - gate_weight: 80 x 16 = 1280 floats = 5KB - vector_weight: 16 x 16 = 256 floats = 1KB - Total: ~26KB This fits in shared memory per block (227KB limit). We could load all weights into shared memory once per block. But with 8 (b,n) per block, each needing to compute independently: - Input loading: each (b,n) loads its own scalar_feat (64 floats) and vector_feat (48 floats) - The norms are computed from vector_feat - Combined is 80 floats (64 scalars + 16 norms) Then do the two linear layers and matrix multiply. Actually, maybe it's cleaner to just have each thread handle its own (b,n) without sharing much. Since the dimensions are small, each thread can keep everything in registers and local memory (which will likely be cached in L1 anyway). Let me try a simpler approach: one thread per (b,n). For B*N = 1024 items, we launch 1024 threads total. Block size of 128, so 8 blocks. Per thread: 1. Load 64 scalar inputs 2. Load 48 vector inputs (16 vectors x 3 dims) 3. Compute 16 norms: norm_i = sqrt(v[i,0]^2 + v[i,1]^2 + v[i,2]^2) 4. For scalar_linear (64 outputs, 80 inputs): For j in 0..63: sum = bias[j] For i in 0..79: sum += combined[i] * W_scalar[i,j] s_out[j] = sum 5. For gate (16 outputs, 80 inputs): For j in 0..15: sum = bias[j] For i in 0..79: sum += combined[i] * W_gate[i,j] gate[j] = sigmoid(sum) 6. For vector transform (16x3 outputs): For i in 0..15: For j in 0..2: sum = 0 For k in 0..15: sum += W_vector[i,k] * v[k,j] v_out[i,j] = gate[i] * sum This is a lot of computation per thread but the dimensions are small enough. Let me count FLOPs: - Norms: 16 * 3 = 48 mults, 16 * 2 = 32 adds, 16 sqrt = 48 FLOPs - Scalar linear: 64 * 80 * 2 = 10240 FLOPs - Gate linear: 16 * 80 * 2 = 2560 FLOPs - Sigmoid: ~10 FLOPs each * 16 = 160 FLOPs - Vector transform: 16 * 3 * 16 * 2 = 1536 FLOPs - Gate multiply: 16 * 3 = 48 FLOPs Total: ~14,592 FLOPs per thread At 67 TFLOPS, this would take 0.2 microseconds if compute-bound. But we're memory-bound. Memory per item: - Read scalar_feat: 64 * 4 = 256 bytes - Read vector_feat: 48 * 4 = 192 bytes - Read weights (if from global): this is shared across all, but still need to load - Write s_out: 64 * 4 = 256 bytes - Write v_out: 48 * 4 = 192 bytes Input/output per item: 256 + 192 + 256 + 192 = 896 bytes For 1024 items: ~917KB of HBM traffic. But weights also need to be loaded: - scalar_weight + bias: ~20.3KB - gate_weight + bias: ~5.1KB - vector_weight: ~1KB Total weights: ~26KB - negligible Total memory: ~943KB. At 3.35 TB/s: ~0.28 microseconds. So this is memory-bound. The key is to avoid temporary allocations and minimize memory traffic. With a fused kernel, we read inputs once, compute everything, write outputs once. This avoids all intermediate tensors (norms, combined, gate unsqueezed, etc.). Now, 1024 threads is not enough to saturate H100. H100 has 132 SMs, each can handle up to 2048 threads. So we want more parallelism. Let me reconsider using multiple threads per (b,n). Option: 32 threads per (b,n), total threads = 1024 * 32 = 32768. Still not saturating but better. Actually, for small matmul-like operations, using warp-level reduction is very efficient. Let me think about this. For scalar_linear with 64 outputs and 80 inputs: - We can have 64 threads each compute one output element. - Each thread loads the 80 inputs (or they are loaded cooperatively into shared memory). - Each thread computes one dot product for one output column. But 64 threads per (b,n) means 1024*64 = 65536 threads total, which is better. For 64 threads per (b,n): - Scalar_linear: 64 threads, 1 output each. Perfect. - Gate: 16 threads compute 1 output each, 48 threads idle. Not great. - Vector transform: 48 outputs, so 48 threads compute 1 each, 16 idle. Also not great. What if we use 16 threads per (b,n)? - Scalar: each thread computes 4 outputs (loop of 4 dot products of length 80) - Gate: each thread computes 1 output (dot product of length 80) - Vector: each thread computes 3 outputs This works! Total work: 1024 * 16 = 16384 threads. Better than 1024 but still light. What if we process multiple (b,n) per thread? Each thread could do several items sequentially. But then we're not increasing parallelism. Actually wait - I should check if 16384 threads is enough. The H100 has 132 SMs * 2048 threads = 270K max concurrent threads. So 16K is about 6% occupancy. But each thread does 14K FLOPs, so maybe it's OK given the small data size. Actually, the problem is that each thread block of 128 threads (8 items) needs its own SM. We have 128 blocks. On 132 SMs, that's almost 1 block per SM. The computation per block is small, so this might be fine. But wait, we need to think about instruction throughput and memory latency hiding. Another option: process multiple outputs in each thread to amortize the input loading. For example, a thread could process 4 (b,n) items sequentially: - Loads weights once (well, they're in shared mem or cached) - For each item, loads inputs, computes, stores outputs - This increases work per thread to ~58K FLOPs But we should first just try a clean implementation and benchmark it. Let me think about using shared memory for inputs. If 16 threads process one (b,n): - They can cooperatively load scalar inputs and vector inputs into shared memory - Then each thread reads from shared memory for its dot products But shared memory is fast. With 8 (b,n) per block of 128 threads, we need 8 * 80 floats for combined = 640 floats = 2.5KB. Plus 8 * 48 floats for vector input = 1.5KB. So about 4KB for inputs in shared memory. Plus ~26KB for weights. Total ~30KB per block. Very reasonable. Actually, let me try a different design. Let's use one warp (32 threads) per (b,n). For scalar_linear (80 -> 64): - Divide 64 outputs among 32 threads: each thread computes 2 outputs - Each thread does 2 dot products of length 80 For gate (80 -> 16): - 16 threads each compute 1 output - Other 16 threads could help with loading or be idle during gate computation For vector MM (16x16 @ 16x3): - 48 elements in output - With 32 threads: not clean division Hmm, what if we use tensor cores? 16x16 is perfect for tensor cores. But with fp32 and small batch, the overhead might not be worth it. Let me try another approach. What about having each block process one (b,n), with enough threads to handle all outputs? Block size = 64 threads for one (b,n): - T0-T63 compute scalar_linear (1 output each) - But wait, for dot products of length 80, we need to load all 80 inputs. - The input loading can be done cooperatively. - Then each thread does a partial dot product reduction? Or each thread loads all 80 inputs itself? If shared memory has the 80 inputs: - Each thread reads 80 floats from shared memory and does 80 multiply-adds for each output it computes. - If each thread computes 1 output, that's 80 FMAs from shared memory. At 2 bytes per float, that's 160 bytes read per thread from shared memory. Shared memory bandwidth is much higher than HBM, so this is efficient. For 64 threads and 64 scalar outputs: - Each thread: 80 FMAs = 160 FLOPs - Total: 10240 FLOPs - good For 16 gate outputs: - 16 threads each do 80 FMAs = 160 FLOPs - Other 48 threads idle during this phase For vector MM (16x16 @ 16x3 = 256 MACs * 3 = 768 FMAs): - We need 48 threads to each produce one element, doing 16 FMAs each - That's only 256 FMAs per thread for the 48 active threads - Other 16 threads idle Total: most threads do 160 + 256 = 416 FMAs which is very little. The memory to read 80 floats from shared memory is done once per thread for all outputs. Wait, each thread only computes one output of scalar_linear. But each thread needs all 80 inputs. So each thread reads 80 floats from shared memory. Total shared memory reads: 64 * 80 * 4 = 20480 bytes. This is fine. For gate, 16 threads each read 80 floats: 16 * 80 * 4 = 5120 bytes. For vector MM, if each of 48 threads does 16 FMAs: each needs a row of A and a column of B. Actually, for C[i,j] = sum_k(A[i,k] * B[k,j]), if thread computes one element, it needs one row of A and one column of B. But B is only 16x3, and A is 16x16. Hmm, for 48 threads we'd have each thread compute one element of the 16x3 output. But there are only 48 elements. So thread (i,j) needs row i of A and column j of B. Since multiple threads need the same row/column, shared memory helps. This is getting complex. Let me think about a simpler approach. ### Simple Approach: One thread per (b,n) Actually, for the given sizes (B=8, N=128), one thread per (b,n) item might be quite efficient because: - The weights are tiny and will stay in L2 cache - Each thread does enough compute (14K FLOPs) to amortize memory access - We avoid synchronization overhead But 1024 threads is not enough for H100. Let me see: 1024 threads / 132 SMs = 8 threads per SM. That's terrible. Hmm, but wait. We have B*N = 1024, and each thread does significant work. Maybe the compiler will vectorize loads/stores and unroll loops. Actually, I realize the key issue: we can process multiple (b,n) items per thread. We could have, say, 256 threads total and each thread does 4 items. Or 128 threads doing 8 items each. But this doesn't increase parallelism. The GPU will just execute fewer threads. Let me think differently. The real bottleneck isn't thread count - it's that each thread block can only use one SM, but if we have few blocks, we underutilize the GPU. For 1024 items, we could: - Block size 64, items per block = 1 → 1024 blocks → saturates all SMs - Block size 128, items per block = 2 → 512 blocks - Block size 256, items per block = 4 → 256 blocks If each block has 64 threads and handles 1 (b,n) item: - 1024 blocks launched - Each block has 64 threads doing ~14K FLOPs - This is much better utilization! But wait, 14K FLOPs per block with 64 threads is ~220 FLOPs per thread. That's very low. Better: block with 64 threads handles 1 (b,n), but the 64 threads cooperate. For scalar_linear with 64 outputs: - Thread i computes output i using a dot product of length 80 + bias - All 64 threads load the 80 inputs cooperatively into shared memory - Then each thread computes its dot product For gate with 16 outputs: - Threads 0-15 compute outputs 0-15 - Threads 16-63 are idle For vector MM: - 48 output elements, how to split among 64 threads? Hmm, for vector MM 16x16 @ 16x3: We can split by output rows (16 rows, each of 3 elements). - Threads 0-15 each compute one row (3 elements each) - Threads 16-63 idle So for each (b,n): - Phase 1: compute norms and load into shared memory - Phase 2: scalar_linear (all 64 threads active) - Phase 3: gate (16 threads active, 48 idle) - Phase 4: vector MM (16 threads active, 48 idle) - Phase 5: apply gate and write results The problem is phases 3-4 have poor thread utilization. What if we compute gate with 64 threads by having each thread compute a partial sum and then do warp shuffle reduction? But 80 inputs / 64 threads = 1.25, not clean. Actually, we can split 80 inputs as: each thread loads ceil(80/64)=2 inputs, does 2 multiplies, then does a warp-level reduction across 64 threads. But 64-thread warp shuffle isn't available (warp is 32 threads). We'd need __shfl_xor across warps, which doesn't work. We'd need shared memory for reduction. So for a 64-thread block, we can do: - Each thread computes partial dot product for its output using a subset of inputs - Store partial sums in shared memory - Reduce within each group of threads assigned to the same output But for scalar_linear, each output has one dedicated thread anyway. So no parallelism there. Let me try yet another design. What if we use 32-thread warps? ### Warp per (b,n) approach Assign 32 threads (1 warp) per (b,n). For scalar_linear (64 outputs): - Each thread computes 2 outputs (64/32 = 2) - Thread t computes outputs t and t+32 - Each thread does 2 dot products of length 80 For gate (16 outputs): - Each thread computes some partial products, then warp shuffle to get final 16 values. - Actually, simpler: 16 threads each compute 1 output directly, 16 threads are idle. For vector MM: - Output is 16x3 = 48 elements - Each thread computes 1 or 2 elements: 32 threads * 1.5 = 48 - Not clean. What if each thread handles 2 outputs of vector MM? 32 * 2 = 64 > 48. So 16 threads handle 1 element, 16 threads handle 2 elements? That's messy. Or have each thread compute a full row of 3 elements? 16 threads each handle one row. But 16 < 32. Hmm. What if we have 16 threads per (b,n)? 16 threads: - Scalar_linear: 64 outputs / 16 = 4 outputs per thread. Each thread does 4 dot products of length 80. - Gate: 16 outputs / 16 = 1 output per thread. Each thread does 1 dot product of length 80. - Vector MM: 48 outputs / 16 = 3 outputs per thread. For vector MM, with 16 threads, a natural split is by output row. Each thread handles one row of the 16x3 output = 3 elements. Thread t computes: v_out[t,0] = sum_k W[t,k] * v[k,0] v_out[t,1] = sum_k W[t,k] * v[k,1] v_out[t,2] = sum_k W[t,k] * v[k,2] Great! Each thread does 3 dot products of length 16. So 16 threads per (b,n) works perfectly: - Thread 0-15 each handle: - 4 scalar outputs - 1 gate output - 1 row of vector output (3 elements) Total threads: 1024 * 16 = 16384 threads. We can use blocks of 128 threads = 8 (b,n) per block. Number of blocks: 1024/8 = 128 blocks. I think this is actually a good design! Let me elaborate: Block size: 128 threads Items per block: 8 (b,n) (each item handled by 16 threads) Number of blocks: ceil(1024 / 8) = 128 blocks Within each block: - Shared memory for weights (since all 8 items in block use same weights): - scalar_weight: 80*64 floats - scalar_bias: 64 floats - gate_weight: 80*16 floats - gate_bias: 16 floats - vector_weight: 16*16 floats - Total: ~26KB - Shared memory for inputs per item: could use registers, or load into shared memory cooperatively - Actually, since 16 threads per item and they all need all 80 inputs, shared memory makes sense - 8 items * 80 inputs = 640 floats = 2.5KB - Plus vector inputs: 8 * 48 = 384 floats = 1.5KB - Total shared memory: ~30KB per block. Very reasonable. Wait, maybe we can load weights once per block into shared memory. But hold on - is loading weights into shared memory worth it? The weights are tiny (26KB total) and read-only. They will be cached in L1/L2. Loading them in shared memory means: 1. Cooperatively load from global to shared 2. Read from shared memory The read from shared memory is slightly faster than L1, but the load cooperation adds synchronization. Given the small size, maybe just reading weights directly from global memory (which will be L2 cached) is fine. Actually, let me think about this more carefully. 128 blocks each loading 26KB of weights = 3.3MB total weight reads. But since all blocks read the same weights, and L2 cache on H100 is large (50MB), the weights will fit in L2 and mostly be cached. So reading weights from global memory is cheap. For inputs: each block reads 8*112 = 896 floats = 3.5KB. This is also tiny. So maybe a simpler kernel with 16 threads per (b,n), reading directly from global memory, is fine. Let me now write out the kernel design more concretely. For each (b,n) assigned to a thread group of 16 threads: Thread `t` in group `g`: - Load scalar_feat[b, n, :] into shared memory or registers (all 16 threads in group cooperate) - Actually, each thread can load some portion - With 16 threads and 64 scalars: each thread loads 4 scalars - Load vector_feat[b, n, :, :] - With 16 threads and 48 floats (16*3): each thread loads 3 floats Then each thread computes its assigned outputs. For scalar_linear outputs assigned to thread t: Output indices: t, t+16, t+32, t+48 (4 outputs) For each output j: sum = bias[j] for i in 0..79: sum += combined[i] * W_scalar[i, j] s_out[b, n, j] = sum For gate output assigned to thread t: Output index: t sum = bias[t] for i in 0..79: sum += combined[i] * W_gate[i, t] gate[t] = 1 / (1 + exp(-sum)) For vector output row t: For j in 0..2: sum = 0 for k in 0..15: sum += W_vector[t, k] * v[k, j] v_out[b, n, t, j] = gate[t] * sum This is quite clean! Let me count memory traffic per thread group (16 threads): Read: - scalar_feat: 64 floats = 256 bytes - vector_feat: 48 floats = 192 bytes - scalar_weight (for 4 outputs each = 64 outputs): but each thread reads cols t, t+16, t+32, t+48, all 80 rows each = 4*80 = 320 floats per thread = 1280 bytes. 16 threads = 20480 bytes total reads from weight. - gate_weight: 1*80 = 80 floats per thread = 320 bytes. 16 threads = 5120 bytes. - vector_weight: 1*16 = 16 floats per thread = 64 bytes. 16 threads = 1024 bytes. Write: - s_out: 64 floats = 256 bytes - v_out: 48 floats = 192 bytes If we assume weights are in L2 (which they should be after first few blocks), the actual HBM traffic is mainly inputs and outputs. Total per-item HBM traffic ≈ 256 + 192 + 256 + 192 = 896 bytes. Wait, but weight reads also count if they miss L2. However, weights are only 26KB total, so once loaded in L2, all subsequent reads are L2 hits. Total execution characteristics: - 1024 items - Each item does ~14.5K FLOPs - Total FLOPs: ~14.8 MFLOPs - Total HBM traffic: ~917KB inputs + outputs + ~26KB weights = ~943KB - Arithmetic intensity: 14.8M / 943K ≈ 15.7 FLOPs/byte At 67 TFLOPS / 3.35 TB/s = 20 FLOPs/byte for compute-boundedness. Since 15.7 < 20, we're memory-bound. The key question is: can we beat PyTorch's multiple-kernel approach by fusing everything? PyTorch reference does: 1. vector_feat.norm(dim=-1): reads 48 floats, writes 16 floats, computes norms 2. torch.cat: reads 64 + 16 floats, writes 80 floats 3. scalar_linear: reads 80 floats, weights, writes 64 floats 4. einsum: reads 48 floats, weights, writes 48 floats 5. vector_gate: reads 80 floats, weights, writes 16 floats 6. sigmoid: reads 16, writes 16 7. unsqueeze + multiply: reads 48 + 16, writes 48 Total PyTorch HBM traffic: - Inputs read once: 256 + 192 = 448 bytes - Outputs written once: 256 + 192 = 448 bytes - But intermediates written and read multiple times Actually, in PyTorch, each op is a separate kernel launching loads/stores. So: 1. norm: read 192, write 64 (for vec_norm) 2. cat: effectively just reorganization, but PyTorch might fuse or allocate new tensor. Actually, it allocates new tensor. read 256 + 64, write 320 (combined) 3. scalar_linear: read 320, read weights, write 256 4. einsum: read 192, read weights, write 192 5. vector_gate: read 320, read weights, write 64 (before sigmoid) 6. sigmoid: read 64, write 64 7. gate multiply: read 192 + 64, write 192 Total reads: 192 + 320 + 256 + 320 + 64 + 192 + 64 = 1408? No, let me be more careful. With eager mode (no fusion): 1. norm kernel: reads vector_feat (192B), writes vec_norm (64B) 2. cat kernel: reads scalar_feat (256B) + vec_norm (64B), writes combined (320B) 3. scalar_linear: reads combined (320B) + weights (20KB) + bias, writes s_out (256B) 4. einsum: reads vector_feat (192B) + W_v (1KB), writes v_out_unscaled (192B) 5. vector_gate: reads combined (320B) + weights (5KB) + bias, writes gate_pre_sigmoid (64B) 6. sigmoid: reads gate_pre_sigmoid (64B), writes gate (64B) 7. unsqueeze+mul: reads gate (64B) + v_out_unscaled (192B), writes v_out (192B) Total HBM traffic (ignoring weights since they're cached in L2): Reads: 192 + 320 + 320 + 192 + 320 + 64 + 256 = 1664 bytes Writes: 64 + 320 + 256 + 192 + 64 + 64 + 192 = 1152 bytes Total: ~2816 bytes per item With our fused kernel: 896 bytes per item That's a 3.14x reduction in memory traffic! So even with imperfect thread utilization, we should get a significant speedup. Let me now think about implementation. ### CUDA Kernel Design ``` // One block handles multiple (B, N) items // Each warp (or group of 16 threads) handles one item // Grid: (num_blocks) where each block has 128 threads = 8 items // Block size: 128 // For thread assignment within block: // item_id = threadIdx.x / 16 // lane_id = threadIdx.x % 16 ``` Let me write the kernel. Actually, let me also think about vector loads. Since scalar_feat is contiguous in the last dimension (64), we can use float4 loads (4 floats at a time). For scalar_feat of 64 floats: 16 float4 loads for 16 threads. Perfect! For vector_feat of 48 floats: not a multiple of 4... 48/16 = 3 floats per thread. Could do float + float2. For weights: - scalar_weight is (80, 64) in PyTorch linear, which is stored transposed for cuBLAS, but in PyTorch Linear, weight is (out_features, in_features) = (64, 80). So W[j, i] gives weight from input i to output j. - When we access W[i, j] with fixed j and varying i, we have strided access. Better to transpose. Actually, in PyTorch nn.Linear, weight is (out_features, in_features), stored row-major. So W[j, i] is the weight connecting input i to output j. For output j, we need all W[j, 0..79]. If thread t accesses W[t, i] for i = 0..79, and W is stored as (64, 80), then each thread accesses a different row, which is contiguous. Good! Similarly, gate_weight is (16, 80), so W_gate[t, i] is contiguous. vector_weight is (16, 16), so W_vector[t, k] is contiguous. So the weight access patterns are actually quite good for coalescing if adjacent threads access adjacent rows. However, each thread in our 16-thread group accesses the same row? No - thread t accesses row t for gate and vector, and rows t, t+16, t+32, t+48 for scalar. So threads 0..15 access different rows. This is actually great for coalescing if multiple items are being processed. Wait, within a warp (32 threads), we'd have: - For item 0: threads 0-15 access rows 0-15 of W_gate - For item 1: threads 16-31 access rows 0-15 of W_gate So if item 0 and item 1 read the same rows of W_gate, the reads from global memory will be broadcast from L2 cache. Since all blocks read the same weights, this should be well cached. For scalar weights, within a warp: - Threads 0-15 access rows 0,16,32,48, then 1,17,33,49, etc. for item 0 - Threads 16-31 access the same pattern for item 1 Actually, each thread loops over i from 0 to 79 and accesses W_scalar[row, i]. Adjacent threads access different rows. So for fixed i, threads t0, t1 access W_scalar[r0, i], W_scalar[r1, i]. These are separated by 80 floats = 320 bytes. So within a 128-byte cache line, there are weights for multiple i values of the same row, but different rows are not in the same cache line. Since different threads access different rows, each thread gets its own cache lines. There might be cache contention but since weights are only 20KB, they're in L2. Hmm, maybe this isn't optimal for coalescing. But for a kernel that's memory-bound on inputs/outputs, weight memory bandwidth is less critical. Let me think about whether we should transpose weights in Python and pass them to the kernel in a more optimal layout. If we transpose scalar_weight to (80, 64), then for output j, W_T[i, j] is the weight. But thread t computes outputs t, t+16, t+32, t+48. So for fixed i, we need W[i, t], W[i, t+16], W[i, t+32], W[i, t+48]. This is scattered access. Worse. Actually, the original layout W[j, i] where j is output and i is input is already good for us because thread t reads row t continuously. For vector_weight (16, 16), the original layout is also fine because thread t reads row t continuously. For gate_weight (16, 80), thread t reads row t continuously. Fine. So PyTorch's default weight layout is already good for our access pattern! Now let me write the actual CUDA code. First, I need to figure out how to compute lane and item indices correctly. ```cuda __global__ void se3_layer_kernel( const float* scalar_feat, const float* vector_feat, const float* scalar_weight, const float* scalar_bias, const float* vector_weight, const float* gate_weight, const float* gate_bias, float* scalar_out, float* vector_out, int B, int N, int scalar_in, int scalar_out_dim, int vector_in, int vector_out_dim ) { const int total_items = B * N; const int threads_per_item = 16; const int items_per_block = blockDim.x / threads_per_item; int item_idx = blockIdx.x * items_per_block + threadIdx.x / threads_per_item; int lane = threadIdx.x % threads_per_item; if (item_idx >= total_items) return; int b = item_idx / N; int n = item_idx % N; // Pointers to this item's data const float* s_in = scalar_feat + item_idx * scalar_in; const float* v_in = vector_feat + item_idx * vector_in * 3; float* s_out_ptr = scalar_out + item_idx * scalar_out_dim; float* v_out_ptr = vector_out + item_idx * vector_out_dim * 3; // Load scalar inputs (64 floats) - each thread loads 4 float s_local[4]; #pragma unroll for (int i = 0; i < 4; i++) { s_local[i] = s_in[lane * 4 + i]; } // Load vector inputs and compute norms float v_local[3]; #pragma unroll for (int i = 0; i < 3; i++) { v_local[i] = v_in[lane * 3 + i]; } float norm = sqrtf(v_local[0]*v_local[0] + v_local[1]*v_local[1] + v_local[2]*v_local[2]); // Now we have in registers: // Thread lane: s_local[0..3] = scalar_in[lane*4 + (0..3)] // v_local[0..2] = vector_in[lane*3 + (0..2)] // norm = norm of the lane-th vector // Compute scalar linear outputs assigned to this thread // Thread lane computes outputs: lane, lane+16, lane+32, lane+48 #pragma unroll for (int out_idx = 0; out_idx < 4; out_idx++) { int j = lane + out_idx * 16; float sum = scalar_bias[j]; // Sum over scalar inputs (64) #pragma unroll for (int i = 0; i < 64; i++) { // Who has scalar input i? Thread i/4, element i%4 float val = __shfl_sync(0xFFFF, s_local[i % 4], i / 4); sum += val * scalar_weight[j * 80 + i]; } // Sum over vector norms (16) #pragma unroll for (int i = 0; i < 16; i++) { float val = __shfl_sync(0xFFFF, norm, i); sum += val * scalar_weight[j * 80 + 64 + i]; } s_out_ptr[j] = sum; } // Compute gate for this thread float gate_sum = gate_bias[lane]; #pragma unroll for (int i = 0; i < 64; i++) { float val = __shfl_sync(0xFFFF, s_local[i % 4], i / 4); gate_sum += val * gate_weight[lane * 80 + i]; } #pragma unroll for (int i = 0; i < 16; i++) { float val = __shfl_sync(0xFFFF, norm, i); gate_sum += val * gate_weight[lane * 80 + 64 + i]; } float gate_val = 1.0f / (1.0f + expf(-gate_sum)); // Compute vector output for this thread's row #pragma unroll for (int j = 0; j < 3; j++) { float sum = 0.0f; #pragma unroll for (int k = 0; k < 16; k++) { float val = __shfl_sync(0xFFFF, v_local[j], k); sum += val * vector_weight[lane * 16 + k]; } v_out_ptr[lane * 3 + j] = gate_val * sum; } } ``` Wait, I used `__shfl_sync` to broadcast values within the 16-thread group. But `__shfl_sync` works within a warp (32 threads), and our 16-thread group might be the lower half of a warp. Since threads 0-15 are in the same warp, `__shfl_sync(0xFFFF, val, src_lane)` works. But we need to make sure the mask covers the right threads. Actually, in CUDA, `__shfl_sync(unsigned mask, T var, int srcLane)` reads from `srcLane` within the warp. The mask must have the calling thread's bit set. If we call with mask=0xFFFF, all threads in the warp participate. In our 16-thread group within a warp of 32, threads 0-15 are one group, threads 16-31 are another. When thread 0 does `__shfl_sync(0xFFFF, val, 5)`, it reads from lane 5 within the warp, which is correct. But there might be a subtle issue: the compiler might not know that `__shfl_sync` is pure, so it might recompute or not optimize well. We should use `#pragma unroll` aggressively. Also, I realize there's a potential issue. The `__shfl_sync` calls in the loops will be executed 64 times and 16 times respectively. The `srcLane` parameter is a compile-time constant in unrolled loops if the loop is fully unrolled and `i` is the induction variable used in a constant expression. Let me check: `__shfl_sync(0xFFFF, s_local[i % 4], i / 4)` - here `i / 4` depends on `i`. If `#pragma unroll` unrolls the loop, then each iteration has a constant `srcLane`. Good. But wait, `s_local[i % 4]` also depends on `i`. If unrolled, the compiler will pick the right element. Actually, there's another concern. `i % 4` and `i / 4` are compile-time constants for each unrolled iteration, but will the compiler inline them? Yes, with `#pragma unroll`. However, `__shfl_sync` takes a variable `var`. For the scalar case, `s_local[i % 4]` isa register variable. For each unrolled iteration, we might be calling `__shfl_sync` with the same `srcLane` but different registers. Wait, I think there's a problem. In the unrolled loop, for a given thread and iteration i: - The thread computes `s_local[i % 4]` (which is its own register value) - Then calls `__shfl_sync` to get the value from the thread at lane `i / 4` - If the calling thread is not the source thread, its own `s_local[i % 4]` isn't used by anyone except itself if it happens to be the source Actually no, `__shfl_sync` returns the value of `var` from the source thread. The `var` parameter is just the local variable of the calling thread. But since we're in a SIMT model, all threads execute the same instruction, and the hardware reads `var` from the source lane. Wait, actually I need to be more careful. `__shfl_sync(mask, var, srcLane)`: - Each thread in the mask provides its `var` - The result for each thread is the `var` from thread `srcLane` - So if all threads call `__shfl_sync(0xFFFF, s_local[0], 0)`, all threads get thread 0's `s_local[0]` This is correct! But there's a question about performance. If we call `__shfl_sync` 64 times for scalar inputs and 16 times for norms, that's 80 shuffle operations per dot product. And we have 4 scalar outputs + 1 gate output = 5 dot products. Total: 400 shuffle operations per thread. Each shuffle is ~1 cycle on modern GPUs, but with 16 threads active per warp (32 threads in warp, but 16 doing work for item A, 16 for item B), the warp executes with both halves active. Actually, since threads 0-15 and 16-31 are doing the same operations but for different items, the warp is fully utilized. Wait, but our group size is 16, and a warp has 32 threads. So two groups fit in one warp. The operations for both groups should execute in parallel. Actually, there might be an issue with `__shfl_sync` within a warp when two groups are using different data. But since they use the same instructions but on different registers, the warp executes both groups together, and `__shfl_sync` within a warp reads from the warp lanes regardless of which group they belong to. For group 0 (threads 0-15), when thread 0 does `__shfl_sync(0xFFFF, val, 5)`, it gets lane 5's value, which belongs to group 0. Good. For group 1 (threads 16-31), when thread 16 does `__shfl_sync(0xFFFF, val, 21)`, it gets lane 21's value, which belongs to group 1. Good. But the mask 0xFFFF includes all 32 threads. The hardware will do the shuffle across all 32 lanes. For group 0, threads 16-31 participate but their results are ignored by threads 0-15. Similarly for group 1. Wait, but we might want to use a narrower mask like `0x0000FFFF` for group 0 and `0xFFFF0000` for group 1. However, `__shfl_sync` only works within a warp, and the mask must include the calling thread. It should be fine to use `0xFFFFFFFF`. Actually, let me reconsider. Instead of shuffle, what if we use shared memory to broadcast the inputs? ### Shared Memory Broadcast Alternative Store inputs in shared memory: - `s_shared[64]` for scalar inputs - `n_shared[16]` for vector norms - `v_shared[16][3]` for vector inputs Each thread reads from shared memory instead of shuffling. For scalar_linear dot product: ```cuda float sum = bias[j]; for (int i = 0; i < 80; i++) { sum += input_shared[i] * W[j * 80 + i]; } ``` Now each thread does contiguous reads from shared memory. No shuffle overhead. Which is faster: shuffle or shared memory? On modern GPUs (Ampere, Hopper), shared memory latency is ~20-30 cycles, while shuffle is ~1-2 cycles. But shuffle requires more instructions and might have overhead if used extensively. Actually, I think shared memory might be better here because: 1. All threads in the group read the same shared memory locations 2. Shared memory broadcasts have the same latency as regular shared memory access 3. We avoid the instruction overhead of many shuffle operations 4. The compiler can better optimize regular memory loads Let me redesign with shared memory instead. For a block of 128 threads = 8 items per block: Shared memory layout: - Combined inputs per item: 80 floats per item - For 8 items: 640 floats = 2560 bytes We could allocate `__shared__ float shared_inputs[8][80];` Each item's 16 threads cooperate to load: - Scalar inputs: threads 0-15 each load 4 floats into shared[0..63] - Vector norm computation: each thread loads its 3 vector components, computes norm, stores norm in shared[64+lane] Wait, but the vector inputs are also needed for the matrix multiply later. So we need either: 1. Store vector inputs in shared memory too 2. Or recompute/reload them If we store vector inputs in shared memory: - `__shared__ float shared_v[8][16][3];` = 8 * 48 * 4 = 1536 bytes Total shared memory: 2560 + 1536 = 4096 bytes = 4KB Plus weights: 26KB Total: 30KB per block Actually, we might not need to store combined inputs if we just store scalar inputs and norms separately. Or even simpler: each thread keeps its own scalar inputs (4 floats) and vector inputs (3 floats) in registers. Then to get another thread's inputs for dot products: - Use `__shfl_sync` to broadcast, OR - Store in shared memory Instead of requiring all threads to have all inputs, what if we restructure the computation to not need broadcasting? For scalar_linear, if thread t computes outputs t, t+16, t+32, t+48: Each output j needs all 80 inputs. So thread t needs to access all 80 inputs for each of its 4 outputs. If inputs are in shared memory, thread t reads 80 * 4 = 320 floats from shared memory. Plus the weights. Wait, with the shuffle approach, thread t's own 4 scalar inputs and 1 norm are in registers. The other inputs come from shuffle. In the shuffle approach, the communication cost is shuffle instructions, not memory loads. With shared memory, thread t reads all inputs from shared memory. Since shared memory is fast, this might be comparable or better. Let me benchmark both approaches mentally: - Shuffle approach: ~400 shuffle ops per thread. Each shuffle is 1-2 cycles, but with dependency chains. - Shared memory approach: ~320 shared memory reads per thread, each ~20 cycles, but many can be overlapped. Hmm, maybe shuffle is actually better for latency hiding because the dot products have plenty of FMAs to overlap with. Actually, wait. In the shuffle approach, each thread only has its own 4 scalar inputs and 1 norm in registers. For the 60 other scalar inputs, it shuffles them in. But each shuffle gives one float from one thread. So for 80 inputs, we need 80 shuffles. In the shared memory approach, thread t reads 80 floats from shared memory, then does 80 FMAs. The FMAs can overlap with the shared memory reads if enough instruction-level parallelism exists. With `#pragma unroll`, the compiler should be able to schedule loads and FMAs efficiently. Let me try the shared memory approach. ```cuda __shared__ float s_shared[ITEMS_PER_BLOCK][80]; __shared__ float v_shared[ITEMS_PER_BLOCK][48]; // or [16][3] // Load scalar inputs for (int i = lane; i < 64; i += 16) { s_shared[item_local][i] = s_in[i]; } // Load vector inputs for (int i = lane; i < 48; i += 16) { v_shared[item_local][i] = v_in[i]; } __syncthreads(); // Compute norms float norms[16]; // Actually norms can be read from v_shared by all threads // But we only need them for dot products // We could precompute norms in shared memory too __shared__ float n_shared[ITEMS_PER_BLOCK][16]; if (lane < 16) { n_shared[item_local][lane] = sqrtf( v_shared[item_local][lane*3+0]*v_shared[item_local][lane*3+0] + v_shared[item_local][lane*3+1]*v_shared[item_local][lane*3+1] + v_shared[item_local][lane*3+2]*v_shared[item_local][lane*3+2] ); } __syncthreads(); // Now thread lane computes its outputs // Scalar linear #pragma unroll for (int out_idx = 0; out_idx < 4; out_idx++) { int j = lane + out_idx * 16; float sum = scalar_bias[j]; float* combined = s_shared[item_local]; #pragma unroll for (int i = 0; i < 80; i++) { sum += combined[i] * scalar_weight[j * 80 + i]; } s_out_ptr[j] = sum; } // Gate float gate_sum = gate_bias[lane]; #pragma unroll for (int i = 0; i < 80; i++) { gate_sum += combined[i] * gate_weight[lane * 80 + i]; } float gate_val = 1.0f / (1.0f + expf(-gate_sum)); // Vector output #pragma unroll for (int j = 0; j < 3; j++) { float sum = 0.0f; #pragma unroll for (int k = 0; k < 16; k++) { sum += v_shared[item_local][k * 3 + j] * vector_weight[lane * 16 + k]; } v_out_ptr[lane * 3 + j] = gate_val * sum; } ``` This looks clean and avoids shuffle overhead. Shared memory usage per block with 8 items: - s_shared: 8 * 80 * 4 = 2560 bytes - v_shared: 8 * 48 * 4 = 1536 bytes - n_shared: 8 * 16 * 4 = 512 bytes - Weights: 26KB Total: ~31KB Wait, do we need n_shared? The norms are in shared memory, but they're just 16 floats. We could compute norms and then use them from shared memory. Or we could avoid n_shared and just read v_shared to compute norms inline when needed. But computing norms 5 times (for 5 dot products) is wasteful. So precomputing norms in shared memory is good. Actually, do we even need s_shared and n_shared separately? We could have combined_shared[80] = [scalar_inputs | norms]. ```cuda __shared__ float combined_shared[ITEMS_PER_BLOCK][80]; // Load scalars for (int i = lane; i < 64; i += 16) { combined_shared[item_local][i] = s_in[i]; } // Load vectors float v_thread[3]; for (int i = 0; i < 3; i++) { v_thread[i] = v_in[lane * 3 + i]; } // Store vectors in shared memory for later use __shared__ float v_shared_local[ITEMS_PER_BLOCK][16][3]; for (int i = 0; i < 3; i++) { v_shared_local[item_local][lane][i] = v_thread[i]; } // Compute and store norm float norm = sqrtf(v_thread[0]*v_thread[0] + v_thread[1]*v_thread[1] + v_thread[2]*v_thread[2]); combined_shared[item_local][64 + lane] = norm; __syncthreads(); ``` n_shared isn't needed as a separate array since norms are the last 16 elements of combined_shared. But wait, we need vectors for the matrix multiply later. So v_shared_local is needed. Total shared: 8 * 80 * 4 + 8 * 48 * 4 = 2560 + 1536 = 4096 bytes for inputs/outputs. Plus weights. Good. Actually, do we need v_shared_local? We each thread already has v_thread[3] in registers. For the matrix multiply: v_out[i,j] = sum_k W[i,k] * v[k,j] Thread lane computes row i=lane. It needs v[k,j] for all k=0..15 and j=0..2. But each thread only has v[lane, 0..2] in registers. It needs other threads' vector values. So either: 1. Store all vectors in shared memory and read from there 2. Use shuffle to get other threads' vector values Since we already need shared memory for combined inputs, and shared memory is fast, storing vectors there too makes sense. Let me reconsider whether we need shared memory or if a fully register-based approach with shuffle is better. In the shuffle approach for vector MM: ```cuda #pragma unroll for (int j = 0; j < 3; j++) { float sum = 0.0f; #pragma unroll for (int k = 0; k < 16; k++) { float val = __shfl_sync(0xFFFF, v_local[j], k); sum += val * vector_weight[lane * 16 + k]; } v_out_ptr[lane * 3 + j] = gate_val * sum; } ``` 16 shuffles per output element, 3 elements = 48 shuffles. Plus 80 shuffles for scalar dot product * 5 = 400 shuffles. Total ~448 shuffles per thread. On Hopper, warp shuffle has ~1 cycle latency and high throughput. So 448 cycles of shuffle overhead. For shared memory approach, we have: - 80 shared reads per dot product * 5 = 400 shared reads - Plus 16 shared reads per matrix multiply row * 3 = 48 shared reads - Total 448 shared reads per thread Shared memory latency on Hopper is ~20-30 cycles. But with FMAs to overlap, effective latency might be hidden. However, with unrolled loops, the compiler might not be able to fully interleave loads and FMAs. Wait, actually in the shared memory approach, we do: ```cuda for (int i = 0; i < 80; i++) { sum += combined[i] * weight[i]; } ``` The `combined[i]` can be loaded early and the FMA can execute while the next load is in flight. With 80 iterations and unrolling, the instruction scheduler should overlap them well. On Ampere/Hopper, LDG (global load) latency is ~300-400 cycles, but LDS (shared load) latency is ~20 cycles. Since each FMA takes ~4 cycles, and we can have many in flight, shared memory loads should be mostly hidden. But wait, if the compiler unrolls the loop, it might generate code that loads multiple values first and then does FMAs, which would be good for latency hiding. But with `#pragma unroll` on a loop of 80, the compiler might generate fully unrolled code with interleaved loads and FMAs. Actually, maybe a better approach is to explicitly pipeline: ```cuda float reg0 = combined[0]; float reg1 = combined[1]; #pragma unroll for (int i = 0; i < 80; i += 2) { sum += reg0 * weight[i]; reg0 = combined[i+2]; // or similar pattern sum += reg1 * weight[i+1]; reg1 = combined[i+3]; } ``` But this is getting complicated and the compiler usually handles it well. Let me try the shared memory approach first because it's cleaner and easier to reason about. If it's not fast enough, we can try shuffle. Actually, I just realized something. If each block has 8 items and we use shared memory for all 8 items' inputs, we need __syncthreads() between loading and computing. But since each item is independent within its 16 threads, and different items don't communicate at all, we could store inputs in shared memory without much overhead. However, there might be bank conflicts. Shared memory has 32 banks on modern GPUs. If threads within a warp access different banks, it's fast. If they access the same bank, it's serialized. For combined_shared[item_local][i]: - If 16 threads in a warp access combined_shared[0][0..15] (assuming item_local=0), each thread accesses a different bank. No conflict. - If they access combined_shared[item_local][same_i], that's a broadcast. No conflict. For weights (in shared memory or global): - Thread t accesses weight[j * 80 + i] where j depends on t. - If j values are separated by multiples of 32 * 4 = 128 bytes, they might hit the same bank. Actually, if weights are in shared memory: - scalar_weight is [64][80] (or actually [64][80] if we think of it as 64 rows of 80 elements each) - Thread t (within a half-warp) accesses rows t, t+16, etc. - Row t starts at t * 80 floats = t * 320 bytes. - Bank = (t * 320 / 4) % 32 = t * 80 % 32. - For t = 0..15: banks = 0, 16, 0, 16, 0, 16... - Hmm, many bank conflicts! Wait, let me recalculate. Each float is 4 bytes, so 32 banks handle 128 bytes. Bank = (address / 4) % 32. Row t starts at address t * 80 * 4 = t * 320 bytes. For t = 0: bank 0 t = 1: 320/4 = 80 % 32 = 16 t = 2: 160 % 32 = 0 t = 3: 240 % 32 = 16 ... So for threads 0,2,4,... they all access bank 0! That's a massive bank conflict for 8 threads. For threads 1,3,5,... they all access bank 16. So if we have 16 threads accessing scalar_weight[j][i] with fixed i and j = lane, lane+16, etc., we get 8-way bank conflicts for bank 0 and 8-way for bank 16. This is bad. However, if weights are in global memory (cached in L1), we don't have shared memory bank conflicts. Bank conflicts only happen for shared memory. So storing weights in shared memory causes bank conflicts! Better to keep weights in global memory (L2 cached) and only use shared memory for the inputs which have better access patterns. For combined_shared[item_local][i]: - All 16 threads in a group access the same item_local but with different i values. - Thread t accesses combined_shared[0][t] for some phase, then [0][t+16], etc. - Address for [0][t]: t * 4 bytes. Bank = t % 32. Threads 0..15 access banks 0..15. No conflicts! - Address for [0][t+16]: (t+16) * 4. Bank = (t+16) % 32. Still no conflicts within 16 threads. So inputs in shared memory have good access patterns. Weights in shared memory have bad patterns. Conclusion: keep weights in global memory (they'll be L2 cached), use shared memory only for inputs/outputs that are shared among threads. Actually, we might not even need shared memory for inputs if we use the shuffle approach. But inputs in shared memory are accessed with no bank conflicts, so that's good. Let me refine the design: ### Refined Design Keep weights in global/L2 memory. Use shared memory for combined inputs and vectors (needed for matrix multiply). Wait, do we need shared memory for vectors? In the matrix multiply, thread lane needs v[k,j] for k=0..15. - If we use shuffle: each thread has v[lane,0..2]. To get v[k,j], thread lane does `__shfl_sync(v_local[j], k)`. - If we use shared memory: thread lane reads `v_shared[item_local][k][j]` from shared memory. Shared memory for vectors: - v_shared[item_local][k][j] has address = k*3 + j elements within the item - Thread lane reads v_shared[0][0][0], v_shared[0][1][0], etc. for j=0 - Addresses: k*12 + j*4 bytes. Bank = (k*3 + j) % 32. - For j=0, k=0..15: banks 0,3,6,9,12,15,18,21,24,27,30,1,4,7,10,13. No conflicts! Great. So shared memory for vectors also has good access patterns. But actually, the vector matrix multiply is: ```cuda for (k = 0; k < 16; k++) { sum += v_shared[k][j] * W[lane][k]; } ``` If threads in a warp execute this together with different `lane` values, they all access the same `v_shared[k][j]` (same k,j) at the same time. This is a broadcast from shared memory, which is optimal! For the weights `W[lane][k]`, each thread accesses a different row (`lane`), but same column `k`. Since weights are in global memory, adjacent threads access different rows, which might not be coalesced. But weights are small and in L2. Hmm, actually `W[lane][k]` with `lane` varying across threads means that for fixed k, threads access W[0][k], W[1][k], etc. If W is stored row-major as [16][16], then W[lane][k] is at address lane*64 + k*4 bytes. For k=0, threads access addresses 0, 64, 128, 192, etc. These are 64 bytes apart. A cache line is 128 bytes, so addresses 0, 64 are in different cache lines. Actually, a cache line holds 32 floats = 128 bytes. So W[0][0..31] is one cache line. W[1][0..31] is another cache line. So each thread accesses a different cache line. With 16 threads accessing 16 different cache lines simultaneously, this is not coalesced. But since these are L2 reads, and multiple threads might hit the same SM's L1 or the shared L2, it might still be fine. But this is a problem because each thread generates a separate L2 request. Similarly for scalar_weight[64][80], for fixed i, threads access rows 0..15 at column i. Each row is 80 floats = 320 bytes. So threads access 0, 320, 640, etc., which are 320 bytes apart. 320/128 = 2.5 cache lines apart. Still not coalesced. This means weight reads from global memory will have poor coalescing. But since weights are tiny and repeated across all blocks, they'll be in L2 and partly in L1. Each SM's L1 can cache a few lines. With 16 threads in a warp accessing 16 different cache lines, the L1 might not be able to satisfy all from local cache; some might need L2 access. However, since the same weights are read by many warps/blocks, they should be well cached in L2. So the actual latency might still be low. But poor coalescing means more memory transactions. For example, with 16 threads accessing 16 different cache lines, we have 16 memory transactions per access pattern instead of 1-2. Can we transpose the weights to fix this? If we store scalar_weight as [80][64] instead of [64][80], then for fixed j, W_T[i][j] has threads accessing W_T[0][j], W_T[1][j], etc. But wait, in our case, each thread needs fixed j and varying i. If W is [64][80], then for fixed j, varying i, threads access different j values, not different i values. Hmm, let me re-examine. In the dot product: ```cuda for (i = 0; i < 80; i++) { sum += combined[i] * W[j][i]; } ``` For a fixed iteration i, threads in a warp have different j values (0..15, 16..31, etc. across the two groups). So for fixed i, threads access W[0][i], W[1][i], W[2][i], etc. These are in adjacent rows. If W is [64][80] row-major, W[0][i] and W[1][i] are 80 floats = 320 bytes apart. So for the scalar linear, within a warp with groups 0 and 1: - Group 0: threads 0-15 access W[0..15][i] - Group 1: threads 16-31 access W[0..15][i] (same rows!) Wait, both groups access the same scalar_weight rows 0..15? No, in group 0, each thread accesses one output. Thread 0 accesses W[0][i] for all i, thread 1 accesses W[1][i], etc. But each thread in group 0 computes 4 outputs. Thread 0 computes outputs 0, 16, 32, 48, so it accesses W[0][i], W[16][i], W[32][i], W[48][i]. For fixed i, within group 0 (threads 0-15): - Thread 0 accesses W[0][i] for its first output - Thread 1 accesses W[1][i] - ... - Thread 15 accesses W[15][i] So for iteration i of the first output, threads 0-15 access rows 0-15. These are 320 bytes apart. Not coalesced. But then for the second output, thread 0 accesses W[16][i], thread 1 accesses W[17][i], etc. The memory access pattern across iterations: thread t sequentially accesses W[t][0], W[t][1], ..., W[t][79] for its first output. This is contiguous! And for its second output: W[t+16][0..79] - also contiguous. So each thread individually has great memory access (contiguous within a row), but threads in a warp don't coalesce their accesses together. In other words, each thread generates contiguous accesses but different threads access different rows. On modern GPUs, this is handled by the memory system. Each warp's memory requests are merged per-cache-line. When threads in a warp access scattered cache lines, each cache line is fetched once. For 16 threads accessing 16 scattered cache lines, we get 16 cache line fetches. Since each cache line is 128 bytes = 32 floats, and each row has 80 floats, a row spans 2.5 cache lines. So 16 rows span 40 cache lines. For each iteration i (single float), we'd need 16 cache lines. But with contiguous access within each thread, adjacent iterations hit the same cache line. So for the first output of each thread, thread t accesses W[t][0..79]. The first access to W[t][0] fetches the cache line containing W[t][0..31]. Then accesses W[t][1..31] are L1 cache hits. Then W[t][32] fetches the next cache line, etc. So each thread gets 3 cache line fetches for its 80 contiguous floats. Total cache line fetches per thread: 3 per output * 4 outputs = 12. For 16 threads: 192 cache line fetches total for scalar weights per item. If we had better coalescing, the 16 threads might have shared cache lines. But since W[t][0..31] and W[t+1][0..31] are on different cache lines, there's no sharing. For L2 cache, since all blocks access the same weights, these cache lines are likely in L2. So each thread fetches from L2 to L1. 192 L2 requests per item. For 1024 items: 196,608 L2 requests. The L2 bandwidth is much higher than HBM, so this should be fine. In fact, the weights are so small and repeatedly accessed that they might stay in L1 cache too, especially since each block processes 8 items and reuses the same weights. For 8 items per block, each block does 8 * 192 = 1536 cache line reads from scalar weights. But the weights are only 64 rows * 80 floats * 4 bytes = 20,480 bytes = 160 cache lines. So the block touches all 160 cache lines about 10 times. Once a cache line is in L1, subsequent accesses are L1 hits. So after the first few iterations, most weight accesses are L1 hits! The initial cold-start misses are insignificant. This means the scalar weight access pattern is actually fine. Similarly for gate_weight (16 rows of 80 floats): 16 * 2.5 = 40 cache lines. Very small. Vector_weight (16x16 = 256 floats = 8 cache lines): tiny. So I think keeping weights in global memory is fine. Now, another idea: instead of 16 threads per item, what if we use 32 threads per item but structure the work differently? With 32 threads per item: - Scalar output: 64 outputs / 32 threads = 2 per thread. Each thread does 2 dot products of 80. - Gate: 16 outputs / 32 threads = 0.5 per thread. Not clean. - Vector MM: 48 outputs / 32 threads = 1.5 per thread. Not clean. What if we assign threads to specific dimensions? 32 threads: - 16 threads for gate and vector (each thread does 1 gate + 1 vector row = 3 elements) - The other 16 threads help with scalar? Not easy to balance. What about 64 threads per item? - 64 scalar outputs / 64 = 1 per thread - 16 gate / 64 = only 16 active - 48 vector / 64 = only 48 active Too much idle time. I think 16 threads per item is the sweet spot. Actually, one more thought: could we have a thread do multiple items sequentially within a single block? Like, 128 threads per block, 16 threads per item, so 8 concurrent items. But each thread could process 2 or more items one after another. For example, 128 threads per block, processing 16 items total (8 concurrently, each thread handles 2 items serially). However, this doubles the work per thread while keeping the same number of blocks. Since we already have 1024 items / 8 = 128 blocks, and a block can only use one SM at a time, reducing blocks to 64 might underutilize the GPU. So 128 blocks with 128 threads each = good occupancy. No need for each thread to do multiple items. Let me now think about warp divergence. With 16 threads per item, and warp size 32, we have 2 items per warp. Threads 0-15 do item A, threads 16-31 do item B. Both groups execute the same instructions on the same clock, just different data. There should be no warp divergence since all 32 threads are active and taking the same code path. Wait, but what about the gate computation? Thread lane in 0..15 computes gate for its lane. All 16 threads in group 0 are active. And all 16 threads in group 1 are active. So the warp has all 32 threads active all the time. Great! One concern: `if (lane < 16)` is not needed since all lanes are 0..15 within each group, but within a warp, lanes range 0..31. However, we don't have any lane-dependent branches. The only difference is that threads 0-15 and 16-31 operate on different memory. So no divergence. Let me now write a complete kernel. Actually, I want to double check something. With 128 threads per block and 16 threads per item, we have 8 items per block. Index calculation: - `item_local = threadIdx.x / 16` (0 to 7) - `lane = threadIdx.x % 16` (0 to 15) For all threads, there are no branches based on lane or item_local (except the guard `if (item_idx >= total_items)`). All threads execute the same instructions, just with different memory addresses and loop bounds. Since the loop bounds are the same for all threads (80 iterations for dot product, etc.), there is no divergence. Wait, for the vector norm: do all threads compute a norm? Yes, each lane loads 3 vector values and computes one norm. All 16 threads in each group do this. All 32 threads in the warp do this. No divergence. OK, I'm confident this design will work well. Let me write the code. One more consideration: the output `v_out` has shape (B, N, vector_out, 3). In PyTorch, this is stored as (B, N, vector_out, 3) = (8, 128, 16, 3). The last dimension (3) is contiguous. So v_out[b, n, i, j] is at index ((b * N + n) * vector_out + i) * 3 + j. Similarly, vector_feat is (B, N, vector_in, 3), so v_in[b, n, k, j] is at ((b * N + n) * vector_in + k) * 3 + j. scalar_feat is (B, N, scalar_in), so s_in[b, n, i] is at (b * N + n) * scalar_in + i. scalar_out is (B, N, scalar_out), so s_out[b, n, j] is at (b * N + n) * scalar_out + j. These are the correct linear indexing formulas. Let me also think about whether we should add `__restrict__` to the pointers. Yes, this helps the compiler optimize. And we should use `const` where appropriate for inputs and weights. Now, the host-side wrapper: ```python def forward_cuda(scalar_feat, vector_feat, scalar_weight, scalar_bias, vector_weight, gate_weight, gate_bias, scalar_out, vector_out): B, N, scalar_in = scalar_feat.shape _, _, vector_in, _ = vector_feat.shape scalar_out_dim = scalar_out.shape[-1] vector_out_dim = vector_out.shape[-2] total_items = B * N threads_per_item = 16 items_per_block = 128 // threads_per_item # 8 num_blocks = (total_items + items_per_block - 1) // items_per_block # Create output tensors # (Actually, outputs are passed in or created in forward) se3_layer_kernel[grid, block]( scalar_feat, vector_feat, scalar_weight, scalar_bias, vector_weight, gate_weight, gate_bias, scalar_out, vector_out, B, N, scalar_in, scalar_out_dim, vector_in, vector_out_dim ) return scalar_out, vector_out ``` Wait, I should create the output tensors in the Python forward and pass them to CUDA, or have CUDA create them? It's more efficient to create them in Python and pass pointers. Let me write the full ModelNew class. We also need to handle the weight and bias extraction from PyTorch layers. The Model class has: - `self.scalar_linear = nn.Linear(scalar_in + vector_in, scalar_out)` - `self.vector_weight = nn.Parameter(torch.randn(vector_out, vector_in))` - `self.vector_gate = nn.Linear(scalar_in + vector_in, vector_out)` In ModelNew __init__, we create the same parameters: - scalar_weight, scalar_bias from scalar_linear - vector_weight - gate_weight, gate_bias from vector_gate During forward, we extract the weights and call the CUDA kernel. Let me also think about whether we need the `scalar_in` parameter in the kernel since combined is always `scalar_in + vector_in`. Actually, in the kernel, we just need the total combined dimension. But it's fine to pass all dimensions for clarity. Here is my planned kernel code: ```cuda #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> #define TILE_M 16 #define ITEMS_PER_BLOCK 8 __global__ void se3_layer_kernel( const float* __restrict__ scalar_feat, const float* __restrict__ vector_feat, const float* __restrict__ scalar_weight, const float* __restrict__ scalar_bias, const float* __restrict__ vector_weight, const float* __restrict__ gate_weight, const float* __restrict__ gate_bias, float* __restrict__ scalar_out, float* __restrict__ vector_out, int B, int N, int scalar_in, int scalar_out_dim, int vector_in, int vector_out_dim ) { const int total_items = B * N; const int threads_per_item = TILE_M; // 16 const int items_per_block = blockDim.x / threads_per_item; int item_idx = blockIdx.x * items_per_block + threadIdx.x / threads_per_item; int lane = threadIdx.x % threads_per_item; if (item_idx >= total_items) return; int b = item_idx / N; int n = item_idx % N; // Input/Output pointers const float* s_in = scalar_feat + item_idx * scalar_in; const float* v_in = vector_feat + item_idx * vector_in * 3; float* s_out_ptr = scalar_out + item_idx * scalar_out_dim; float* v_out_ptr = vector_out + item_idx * vector_out_dim * 3; // Shared memory for combined inputs and vector features // Layout: [items_per_block][scalar_in + vector_in] for combined // [items_per_block][vector_in * 3] for vectors extern __shared__ float shared_mem[]; float* combined_shared = shared_mem; // size: items_per_block * (scalar_in + vector_in) float* v_shared = shared_mem + items_per_block * (scalar_in + vector_in); // size: items_per_block * vector_in * 3 int item_local = threadIdx.x / threads_per_item; int combined_offset = item_local * (scalar_in + vector_in); int v_offset = item_local * vector_in * 3; // Load scalar inputs cooperatively // scalar_in = 64, threads_per_item = 16, so each thread loads 4 #pragma unroll for (int i = lane; i < scalar_in; i += threads_per_item) { combined_shared[combined_offset + i] = s_in[i]; } // Load vector inputs cooperatively // vector_in * 3 = 48, threads_per_item = 16, so each thread loads 3 #pragma unroll for (int i = lane; i < vector_in * 3; i += threads_per_item) { v_shared[v_offset + i] = v_in[i]; } __syncthreads(); // Compute vector norms and store in combined_shared[scalar_in + lane] float v0 = v_shared[v_offset + lane * 3 + 0]; float v1 = v_shared[v_offset + lane * 3 + 1]; float v2 = v_shared[v_offset + lane * 3 + 2]; float norm = sqrtf(v0 * v0 + v1 * v1 + v2 * v2); combined_shared[combined_offset + scalar_in + lane] = norm; __syncthreads(); // Scalar linear: thread lane computes outputs lane, lane+16, lane+32, lane+48 #pragma unroll for (int out_idx = 0; out_idx < (scalar_out_dim / threads_per_item); out_idx++) { int j = lane + out_idx * threads_per_item; float sum = scalar_bias[j]; #pragma unroll for (int i = 0; i < scalar_in + vector_in; i++) { sum += combined_shared[combined_offset + i] * scalar_weight[j * (scalar_in + vector_in) + i]; } s_out_ptr[j] = sum; } // Gate linear: thread lane computes output lane float gate_sum = gate_bias[lane]; #pragma unroll for (int i = 0; i < scalar_in + vector_in; i++) { gate_sum += combined_shared[combined_offset + i] * gate_weight[lane * (scalar_in + vector_in) + i]; } float gate_val = 1.0f / (1.0f + expf(-gate_sum)); // Vector transform: thread lane computes row lane (3 outputs) #pragma unroll for (int j = 0; j < 3; j++) { float sum = 0.0f; #pragma unroll for (int k = 0; k < vector_in; k++) { sum += v_shared[v_offset + k * 3 + j] * vector_weight[lane * vector_in + k]; } v_out_ptr[lane * 3 + j] = gate_val * sum; } } ``` Wait, there are a few issues: 1. `extern __shared__ float shared_mem[];` requires dynamic shared memory allocation. We need to compute the size in Python when launching the kernel. 2. With `#pragma unroll` on the loops, but loop bounds are runtime variables (scalar_in, vector_in, etc.), the compiler might not know to unroll them. For the specific sizes (64, 16, 16), we should either use template parameters or set the sizes as constants. Actually, we can template the kernel with fixed sizes! Or we can just use constant values within the kernel since the problem specifies fixed sizes. But ModelNew should accept the same constructor arguments, so it needs to work with arbitrary sizes. However, the benchmark tests with specific sizes. We can either: - Make the kernel generic (slower due to non-constant loop bounds) - Make the kernel templated with compile-time sizes - Or use if/constexpr inside the kernel Given that the test uses specific sizes, and compile-time bounds allow better optimization, let me use a templated kernel if possible. But with `load_inline`, we can pass template parameters. Actually, an easier approach: since the problem states `scalar_in = 64`, `scalar_out = 64`, `vector_in = 16`, `vector_out = 16`, we can hardcode these in the kernel. But ModelNew should accept constructor arguments. So the kernel can have the dimensions as template parameters. In load_inline, we'd compile a kernel with these specific template instantiations. But to make it generic, we can compile multiple versions or use runtime parameters. Actually, for speed, hardcoding the sizes in the kernel is best. We can write the kernel with template parameters `SI, SO, VI, VO` and instantiate them in the host wrapper. Let me check if `load_inline` supports C++ templates. Yes, it just compiles the CUDA code as C++, and we can write templated kernels. The host function can call the appropriate template instantiation. But actually, the simplest approach for this problem is to just write the kernel with runtime bounds but use `#pragma unroll` somehow. With runtime bounds, `#pragma unroll` only unrolls if the compiler can determine a small constant upper bound. Since 64 and 16 are passed as parameters, the compiler doesn't know they're small constants. However, even without unrolling, the kernel should still be quite fast. The loops are short (max 80 iterations), so the loop overhead is minimal. But to get best performance, let's use template parameters or at least compile-time constants in the kernel. We can do this: ```cuda template <int SI, int SO, int VI, int VO> __global__ void se3_layer_kernel(...) { ... } ``` And in the host function: ```cpp if (scalar_in == 64 && scalar_out == 64 && vector_in == 16 && vector_out == 16) { se3_layer_kernel<64, 64, 16, 16><<<grid, block, shared_mem_size>>>(...); } ``` But wait, the `load_inline` function mapping means we need to have a non-templated host function in our C++ source that Python calls. That host function can then call the templated kernel. Let me re-read the problem: "Do not write your own `PYBIND11_MODULE(...)` block in `cpp_sources`. `load_inline` auto-generates one from the `functions=[...]` argument; including your own causes a duplicate-symbol redefinition error." So the host function we expose to Python is: ```cpp std::vector<torch::Tensor> se3_layer_forward( torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_weight, ... ) { // Allocate outputs // Call kernel // Return outputs } ``` This host function can call a templated kernel internally. Let me write the full C++ and Python code. Actually, let me double check: `scalar_weight` from `nn.Linear(80, 64)` has shape (64, 80). When we access `scalar_weight[j * 80 + i]`, that's assuming the weight is stored in a flat buffer of 64*80 = 5120 floats, with row-major layout. PyTorch tensors are row-major by default, so `scalar_weight[j, i]` is at `scalar_weight.data_ptr<float>()[j * 80 + i]`. This is correct. Similarly, `gate_weight` is (16, 80), so `gate_weight[lane * 80 + i]` accesses row `lane`, column `i`. Correct. `vector_weight` from `nn.Parameter(torch.randn(16, 16))` is (16, 16). `vector_weight[lane * 16 + k]` accesses row `lane`, column `k`. Correct. One thing about bias: `scalar_bias` is (64,), `gate_bias` is (16,). For the combined inputs, `scalar_in + vector_in = 64 + 16 = 80`. Now, an important question: vector norm computation. PyTorch's `.norm(dim=-1)` computes the L2 norm: `sqrt(sum(x_i^2))`. This is what I'm doing. And the sigmoid gate: PyTorch's `sigmoid(x)` is `1 / (1 + exp(-x))`. This is what I'm doing. The vector transformation: `einsum("oi,bnid->bnod")` where o=vector_out=16, i=vector_in=16, d=3. C[b,n,o,d] = sum_i W[o,i] * V[b,n,i,d] For fixed b,n,o, and d=0..2: v_out[o, d] = sum_i W[o,i] * v[i,d] This is exactly what I'm computing. Gate multiplication: `gate.unsqueeze(-1) * v_out` gate has shape (B, N, vector_out=16), unsqueezed to (B, N, 16, 1). v_out has shape (B, N, 16, 3). So the multiplication broadcasts: gate[b,n,o,0] * v_out[b,n,o,d] for d=0,1,2. This is what I'm doing with `gate_val * sum`. And the concatenation: `torch.cat([scalar_feat, vec_norm], dim=-1)` gives combined[..., 0:64] = scalar_feat, combined[..., 64:80] = vec_norm. This is what I do. OK, mathematically correct. One optimization: I should use `__ldg()` for read-only global memory accesses. Or at least add `__restrict__ const` qualifiers. Let me add `__restrict__`. Actually, `__restrict__` on all pointers is good. The compiler can assume no aliasing. Let me also check shared memory size calculation: - combined_shared: items_per_block * (scalar_in + vector_in) floats - v_shared: items_per_block * vector_in * 3 floats Total shared floats: items_per_block * (SI + VI + 3*VI) = items_per_block * (SI + 4*VI) With SI=64, VI=16: 64 + 64 = 128 floats per item With 8 items: 1024 floats = 4KB So shared memory is 4KB per block. Very small, plenty of room. Wait, I also need to consider whether to use float4 loads for vector_feat and scalar_feat. For scalar_feat: 64 floats / 16 threads = 4 floats per thread. We can use `float4` loads: ```cuda const float4* s_in4 = (const float4*)s_in; float4 val = s_in4[lane]; // val.x, val.y, val.z, val.w are the 4 floats ``` For vector_feat: 48 floats / 16 threads = 3 floats per thread. Not a multiple of 4, so float4 doesn't work cleanly. We could use a `float4` and ignore one component, but that might read out of bounds. Or just do 3 scalar loads. Anyway, using vector loads helps coalesce memory access. Let me use float4 for scalar loads. Actually, for simplicity and since the memory bandwidth is already low, scalar loads might be fine. But let me try vector loads for potentially better performance. Hmm, but then we need to store the loaded values in shared memory. Writing 4 floats from a float4 to shared memory is straightforward. Let me think about whether this is worth it. The HBM bandwidth is very high (3.35 TB/s), and each warp reads 64 floats = 256 bytes for scalars and 48 floats = 192 bytes for vectors. At 3.35 TB/s, a warp can read 448 bytes in ~0.13 nanoseconds? No, that's wrong. Wait, let's compute. HBM bandwidth 3.35 TB/s = 3.35 * 10^12 bytes/s. Per warp: 448 bytes. Time = 448 / (3.35 * 10^12) * 10^9 = 0.134 microseconds for one warp? Hmm no, that's if a single warp monopolizes the bus. Actually, multiple warps are active. The read of 448 bytes per warp happens. With 32 threads, an uncoalesced read of 64 floats might take multiple transactions. But with coalesced float4 loads, we might get fewer transactions. For scalar_feat with 16 threads each doing float4: - Thread 0 reads address 0, thread 1 reads address 16, etc.? No, thread t reads address t*16. - Actually, if each thread loads a float4 at offset `lane * 4`, then thread 0 reads bytes 0-15, thread 1 reads bytes 16-31, etc. - These are perfectly contiguous and adjacent. The warp issues a single memory transaction for 16*16 = 256 bytes. Great! For vector_feat with 16 threads each loading 3 floats: - Thread 0 reads bytes 0-11, thread 1 reads bytes 12-23, etc. - These are perfectly contiguous. A single transaction for 16*12 = 192 bytes. Good! So even with scalar loads (not float4), the loads should be coalesced because the threads access contiguous memory. Actually, let me think about this more carefully. If 16 threads each load a scalar float at addresses `lane * 1` (in units of floats), then thread 0 reads address 0, thread 1 reads address 4 bytes, thread 2 reads address 8 bytes, etc. This is contiguous and perfectly coalesced into a single memory transaction. For a warp on Ampere/Hopper, a warp can load 128 bytes per transaction. So 16 scalar loads of 4 bytes each = 64 bytes, which fits in one transaction. But actually, for 32 threads in a warp, all threads need to participate or some might not. In our case, threads 0-15 load item A, threads 16-31 load item B. For threads 0-15: addresses 0, 4, 8, ..., 60 = 64 bytes total. This fits in one 128-byte transaction. For threads 16-31: addresses 64, 68, 72, ..., 124 = 64 bytes total. Also fits in one 128-byte transaction. Wait, but these are for different items! Threads 0-15 and 16-31 are accessing different items' scalar features. If item A is at base address X, and item B is at base address Y, then threads 0-15 access X+0 to X+60, and threads 16-31 access Y+0 to Y+60. The memory controller merges requests from the same warp. So it may issue two transactions: one for X to X+127, and one for Y to Y+127. This is fine. The loads are still coalesced. So scalar loads are fine. No need for float4 unless we want to reduce instruction count slightly. Let me proceed with scalar loads for simplicity. Now, let me write the full Python + CUDA code. Actually, one more optimization idea. Instead of using shared memory for combined inputs, what if we use register shuffling? With shuffle, each thread keeps its own 4 scalar inputs and 1 norm in registers. For dot products, it gets other threads' values via shuffle. Let me count instructions: For scalar dot product of length 80: - 80 iterations - Each iteration: `__shfl_sync` + FMA - `__shfl_sync` is one instruction on modern GPUs For 4 dot products: 320 shuffles + 320 FMAs. For gate: 80 shuffles + 80 FMAs. For vector MM: - For each of 3 output elements: 16 shuffles + 16 FMAs - Total: 48 shuffles + 48 FMAs. Total per thread: 448 shuffles + 448 FMAs. On Hopper, FMA throughput is very high (can execute multiple per cycle per SM), and shuffle throughput is also high. The total instruction count is quite low. Compare with shared memory approach: For scalar dot product: - 80 shared loads + 80 FMAs per output - For 4 outputs: 320 loads + 320 FMAs - For gate: 80 loads + 80 FMAs - For vector: 48 loads + 48 FMAs Total: 448 loads + 448 FMAs Shared loads have latency ~20 cycles, while shuffle has latency ~1 cycle. But since FMAs and loads can be interleaved, and with loop unrolling, the effective throughput might be similar. However, shared memory has a potential advantage: all threads read the same shared memory location in each iteration (broadcast), which is very fast. But on Hopper, shared memory broadcast might also take a few cycles. Actually, on Ampere/Hopper, `__shfl_sync` is so fast that the shuffle approach might be better. It also avoids `__syncthreads()` calls and shared memory setup. Let me try the shuffle approach to avoid shared memory bank issues and synchronizations. Wait, I realize the shuffle approach has a problem with the "items per warp" setup. A warp has 32 threads. If we have 2 items per warp (16 threads each), then when threads 0-15 do `__shfl_sync(0xFFFF, val, k)`, they get values from lanes 0-15. But the mask 0xFFFF includes all 32 threads. The result for thread 0 is the value of `val` from lane `k` (which is in group 0). This is fine because lane k has the right item's data. Similarly for group 1 (threads 16-31), `__shfl_sync(0xFFFF, val, k)` gets from lane k, which is in group 0, not group 1! Oh no! This is a critical bug. `__shfl_sync` operates on lanes 0-31 of a warp. If threads 16-31 call `__shfl_sync(0xFFFF, val, 5)`, they get the value from lane 5, which belongs to item A, not item B. So for group 1, we need to add an offset. Thread 16 + lane in group 1 should call `__shfl_sync(0xFFFF, val, 16 + k)`. Since lane in group 0 is `threadIdx.x % 16`, and the warp lane is also `threadIdx.x % 32`, group 0 occupies lanes 0-15 and group 1 occupies lanes 16-31. So for group 0: `src_lane = k` (where k in 0..15) For group 1: `src_lane = 16 + k` But all 32 threads in the warp execute the same instruction. So `src_lane` must be warp-lane, meaning for group 0 it's `k`, for group 1 it's `16 + k`. In the code, we can compute: ```cuda int warp_lane = threadIdx.x % 32; int group_base = (warp_lane / 16) * 16; // Then src_lane = group_base + k ``` Since `group_base` is 0 for group 0 and 16 for group 1, and all threads in the warp compute it the same way? No, threads 0-15 get group_base=0, threads 16-31 get group_base=16. But this creates divergence only in the `src_lane` parameter, not in control flow. The `__shfl_sync` instruction takes a runtime `srcLane` parameter, so all threads in the warp execute the same instruction with the same `srcLane`. But since `group_base` is different for the two halves, this would be a problem if they tried to execute different instructions. However, `__shfl_sync` is a single instruction; each thread uses its own `srcLane` value (because `srcLane` is a register value). Wait, is `srcLane` in `__shfl_sync` a compile-time constant or runtime variable? `__shfl_sync(unsigned mask, T var, int srcLane, int width=warpSize)` On modern NVIDIA GPUs, `srcLane` can be a runtime variable. Each thread in the warp can use a different `srcLane` value. The hardware supports this. But there's a catch: if different threads in the warp have different `srcLane` values, the execution is still uniform (no divergence), but the hardware might need to do more work. Actually, for `srcLane` as a runtime variable, the warp still executes one `SHFL` instruction, but with per-thread `srcLane` registers. This is supported. So the shuffle approach can work with the correct lane calculation. Let me write the shuffle-based kernel: ```cuda __global__ void se3_layer_kernel(...) { // ... index setup ... int warp_lane = threadIdx.x % 32; // 0..31 within warp int group_base = (warp_lane / 16) * 16; int lane_in_group = warp_lane % 16; // Load scalar inputs: each thread loads 4 float s_local[4]; #pragma unroll for (int i = 0; i < 4; i++) { s_local[i] = s_in[lane_in_group * 4 + i]; } // Load vector inputs: each thread loads 3 float v_local[3]; #pragma unroll for (int i = 0; i < 3; i++) { v_local[i] = v_in[lane_in_group * 3 + i]; } // Compute norm float norm = sqrtf(v_local[0]*v_local[0] + v_local[1]*v_local[1] + v_local[2]*v_local[2]); // Compute scalar linear outputs #pragma unroll for (int out_idx = 0; out_idx < 4; out_idx++) { int j = lane_in_group + out_idx * 16; float sum = scalar_bias[j]; #pragma unroll for (int i = 0; i < 64; i++) { int src_thread = group_base + i / 4; int src_reg = i % 4; float val = __shfl_sync(0xFFFFFFFF, s_local[src_reg], src_thread); sum += val * scalar_weight[j * 80 + i]; } #pragma unroll for (int i = 0; i < 16; i++) { int src_thread = group_base + i; float val = __shfl_sync(0xFFFFFFFF, norm, src_thread); sum += val * scalar_weight[j * 80 + 64 + i]; } s_out_ptr[j] = sum; } // Gate float gate_sum = gate_bias[lane_in_group]; #pragma unroll for (int i = 0; i < 64; i++) { int src_thread = group_base + i / 4; int src_reg = i % 4; float val = __shfl_sync(0xFFFFFFFF, s_local[src_reg], src_thread); gate_sum += val * gate_weight[lane_in_group * 80 + i]; } #pragma unroll for (int i = 0; i < 16; i++) { int src_thread = group_base + i; float val = __shfl_sync(0xFFFFFFFF, norm, src_thread); gate_sum += val * gate_weight[lane_in_group * 80 + 64 + i]; } float gate_val = 1.0f / (1.0f + expf(-gate_sum)); // Vector transform #pragma unroll for (int j = 0; j < 3; j++) { float sum = 0.0f; #pragma unroll for (int k = 0; k < 16; k++) { int src_thread = group_base + k; float val = __shfl_sync(0xFFFFFFFF, v_local[j], src_thread); sum += val * vector_weight[lane_in_group * 16 + k]; } v_out_ptr[lane_in_group * 3 + j] = gate_val * sum; } } ``` This should work correctly. But there are a lot of `__shfl_sync` calls. Let me think if this is actually efficient. The key question is: can the compiler optimize this? With `#pragma unroll`, yes, the loops get unrolled. Each iteration becomes a shuffle + FMA. The instruction scheduler can overlap FMAs with shuffles. However, 400 shuffle instructions per thread is a lot. Pipe utilization for shuffle is high but not infinite. Let me check if there's a better way. Actually, I realize there's a simpler approach: use a warp-level matrix multiplication or use shared memory more efficiently. What if we assign one thread to one (b,n) and just compute everything serially? Then we have 1024 threads, each with ~14K FLOPs. On H100, each SM can handle 2048 threads. With 1024 threads total and 132 SMs, we have ~8 threads per SM. But each thread does 14K FLOPs. At ~100 GFLOPS per thread? No, that's not how it works. The issue with too few threads is that memory latency can't be hidden. When a thread stalls on a memory load, there might not be enough other threads to keep the SM busy. But each SM on H100 can handle many warps. With 8 threads per SM, that's only 1/4 of a warp. The SM will be mostly idle, waiting for this tiny warp to finish. So we definitely need more threads. The 16-thread-per-item approach gives us 16K threads total, which is better. With 128 threads per block and 128 blocks, we should have decent utilization. Actually, I just realized: with the shared memory approach and 128 threads per block (8 items), we have 128 blocks. H100 has 132 SMs, so almost 1 block per SM. Each block has 128 threads. But each SM can handle up to 2048 threads. So only ~6% thread occupancy per SM. However, occupancy (number of resident threads / max threads) is less important than having enough warps to hide latency. With 128 threads per block = 4 warps per block. If each SM has 1 block, that's 4 warps per SM. Ampere/Hopper needs at least a few warps to hide latency. But the workload per warp is quite high. Let me estimate time: Per warp (32 threads, handling 2 items): - Each item: 14.5K FLOPs - 2 items: 29K FLOPs - Each warp: 29K FLOPs / 32 threads ≈ 906 FLOPs per thread? No, total warp FLOPs = 29K. Actually, let me think about warp execution time. A warp executes 32 threads in parallel. For our kernel, a warp handles 2 items with the shared memory approach. The key operations in the warp: - Memory loads: 2 items * (256 + 192 bytes) = 896 bytes per warp - Shared memory stores and loads - Compute: ~29K FLOPs At peak FMA throughput, each SM can do ~200 GFLOPS (actually H100 does 67 TFLOPS for FP32 across 132 SMs, so ~500 GFLOPS per SM). At 500 GFLOPS, 29K FLOPs take ~58 nanoseconds per SM. But with only 4 warps per SM, can we fully utilize the FMA units? Probably not. Each warp has limited instruction-level parallelism. But with loop unrolling and many independent FMAs, the compiler might be able to keep the pipelines busy. Hmm, actually the shared memory loads might be the bottleneck. Let me reconsider: maybe 64 threads per item? That would give us 1024*64 = 65536 threads total. Or even 128 threads per item = 131072 threads. This would give much better occupancy. With 64 threads per item: - 64 scalar outputs: 1 per thread, each thread does a dot product of length 80 - 16 gate outputs: 16 threads compute 1 each, 48 threads idle - 48 vector outputs: 48 threads compute 1 each, 16 threads idle The idle threads are a problem. Total active threads per item = 64 + 16 + 48 = 128, but not simultaneously. At any given phase, only 64 threads are active. What if we interleave the computation? Like, while threads 0-63 do scalar linear, threads 0-15 also prepare for gate? No, they can't do both at once. What about having 32 threads per item and using warp shuffles for reductions? For example, for scalar_linear with 64 outputs and 80 inputs: - 32 threads, each computes 2 outputs (64/32 = 2) - Each output is a dot product of 80 inputs - Thread t computes outputs t and t+32 But we still have the gate issue: 16 outputs with 32 threads, so we could have all 32 threads participate by having each thread compute a partial sum and then reduce. For dot product of length 80 with 32 threads: - Each thread computes partial sum of 80/32 = 2.5 elements. Not clean. - Or 16 threads compute 5 elements each, then 16 threads are idle during partial computation but then participate in warp shuffle reduction? Hmm, this is getting complicated. Let me look for a simpler approach. ### Alternative: Use cuBLAS/cuDNN? Actually, cuBLAS is very optimized for matrix multiplication. For the linear layers, we could call cuBLAS. But the overhead of multiple kernel launches and intermediate allocations would likely make it slower than a fused kernel. What about using PyTorch's custom CUDA kernel with `torch.compile`? No, we need a raw CUDA kernel. ### Another Design: One Warp = One (B,N), 32 Threads Let me try to make 32 threads per item work. Scalar output: 64 outputs. - 32 threads, each computes 2 outputs. - Each thread does 2 dot products of length 80. Gate output: 16 outputs. - We could have 32 threads each compute a partial dot product, then warp shuffle reduce. - 80 inputs / 32 threads = 2.5. Not clean. - Or we could load the 80 inputs into shared memory (or via shuffle broadcast) and have 16 threads compute the 16 outputs, while 16 threads are idle. - Then for vector MM, have 16 threads each compute 3 outputs (1 row), while 16 threads are idle. So for phases with 32 threads, 16 threads are idle. That's 50% utilization during those phases. But the scalar phase uses all 32 threads. If scalar phase is the longest (4 dot products * 80 = 320 FMAs per thread), and gate is 1 dot product * 80 = 80 FMAs, and vector MM is 3 dot products * 16 = 48 FMAs, then the scalar phase dominates. So 32 threads per item gives us 1024*32 = 32768 threads. Blocks of 128 threads = 4 items per block. 256 blocks total. With 256 blocks, we have ~2 blocks per SM. That's better than 1 block per SM! But the thread occupancy is still low: 128 threads per block, 2 blocks per SM = 256 threads per SM out of 2048 max = 12.5% occupancy. For latency hiding, 256 threads per SM = 8 warps. This is typically enough for some latency hiding, though not ideal. But compared to 16 threads per item (128 blocks = ~1 block per SM = 128 threads per SM = 4 warps = 6.25% occupancy), 32 threads per item is better. What about 64 threads per item? 1024*64 = 65536 threads. Blocks of 128 threads = 2 items per block. 512 blocks total = ~4 blocks per SM. 256 threads per SM = 8 warps. Same occupancy as 32 threads/item. But 64 threads/item means in scalar phase, 64 threads active; in gate phase, only 16 active (25%); in vector phase, 48 active (75%). Average utilization lower than 32 threads/item. So 32 threads/item seems like a sweet spot for occupancy. But the idle threads during gate and vector phases are a concern. Actually, wait. Can we restructure the computation so that all 32 threads are productive all the time? What if we combine gate and vector computation? For vector MM, thread t computes some rows, but each row needs the corresponding gate value. Gate is a scalar linear: `gate = sigmoid(W_gate @ combined + b_gate)`. This is a mat-vec of size 80 -> 16. What if we break the mat-vec into partial sums computed by all 32 threads and then reduce? For the gate calculation: C[j] = sum_i(W[j,i] * x[i]) + b[j] With 32 threads and j=0..15: Option 1: 16 threads each compute one j. Other 16 threads idle. Option 2: All 32 threads compute partial sums for all j, then reduce. For partial sums: - Each thread computes a partial dot product for all 16 outputs. - Thread t sums over inputs i = 2t, 2t+1 (if we split 80 inputs / 32 = 2.5, not clean) - Or thread t sums over a subset of inputs for all outputs. - After partial sums, we need to reduce across the warp for each output. With 32 threads, we can have each thread handle 80/32 = 2.5 inputs. That's messy. But what if we pad combined to 96 inputs? Then 96/32 = 3 inputs per thread. But combined is actually 80, and padding wastes compute. Or we could have 16 threads compute 5 inputs each (total 80) for all 16 outputs, and 16 threads do nothing. Then do a reduction. Actually, for Hopper, warp shuffle reduction is very fast. Let me think about this. If 16 threads each compute a partial sum of 5 inputs for all 16 gate outputs: Each thread computes for j=0..15: partial[j] = sum_{i in my_subset} W[j,i] * x[i] Then do a warp-level reduction across 16 threads for each j. But this requires storing 16 partials per thread in registers, then doing 16 parallel reductions. Warp shuffle can reduce a value across all active threads. But we need to reduce 16 different values. Actually, we could use the `__reduce_add_sync` or `__ballot_sync` type operations, but for float sums across specific threads, we typically use sequential shuffles. For 16 values per thread, reducing them across 16 threads would require 16 * log2(16) = 64 shuffle+add operations. That's a lot. Hmm, this doesn't seem efficient. Let me just go with the 16-thread-per-item design. It has clean thread assignment and all threads are always active. The only downside is lower occupancy, but 16K threads should still give decent performance. Actually, thinking more about occupancy: H100 has 132 SMs, each with 2048 max threads. That's 270K total threads. With 16K threads, we're at 6% of max occupancy. Is 6% occupancy enough? For memory-bound kernels, high occupancy is crucial for latency hiding. But our kernel is quite compute-intensive per thread (14K FLOPs per item with 16 threads = ~900 FLOPs per thread). At 500 GFLOPS per SM and 128 threads per SM, that's about 1/4 of peak FLOPS per SM. So each thread executes for a few microseconds. The memory loads happen at the beginning and are cached. So memory latency is less of an issue than in a typical memory-bound kernel. I think 16 threads per item is fine for this workload size. One more idea: we can increase block size. Instead of 128 threads per block, use 256 threads per block = 16 items per block. Total blocks: 1024/16 = 64 blocks. Per SM: ~0.5 blocks. That's even worse for occupancy! Wait, actually the number of blocks per SM depends on block size. A larger block size means fewer blocks can fit on an SM simultaneously. But the total thread count is the same. With 256 threads per block and 64 blocks total: - Total threads: 64 * 256 = 16384 - Threads per SM: 16384 / 132 ≈ 124 threads per SM - Warps per SM: ~4 With 128 threads per block and 128 blocks total: - Total threads: 128 * 128 = 16384 - Threads per SM: ~124 - Warps per SM: ~4 Same thread count! So block size doesn't change thread occupancy. But larger blocks might have better cache utilization because more threads share L1 cache and more work per block means the weights stay in L1 longer. With 256 threads per block and 16 items per block: - Weight accesses: 16 items * same weights. The weights are accessed repeatedly within the block, staying hot in L1. Actually, all blocks access the same weights. Whether the weights stay in L1 depends on L1 size per SM. H100 has ~128KB L1 per SM (maybe more). Our weights are ~26KB. So multiple blocks' weight working sets might fit in L1 if they don't interfere. But with many blocks running on the same SM (if occupancy were higher), weight evictions from L1 would happen. With our low occupancy, L1 is mostly empty, so weights can stay there. I think block size of 128 or 256 won't make a huge difference. Let me just use 128 threads. Now, there's another important consideration: the evaluation system runs with a specific batch size and might warm up the GPU. The reference implementation is a PyTorch nn.Module. Our ModelNew should be faster. To maximize speedup, I should also minimize Python overhead. ModelNew.forward should be lean. Let me write the code now. ### Full Python + CUDA Code ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline cuda_source = """ #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> #define TILE_M 16 template <int SI, int SO, int VI, int VO> __global__ void se3_layer_kernel( const float* __restrict__ scalar_feat, const float* __restrict__ vector_feat, const float* __restrict__ scalar_weight, const float* __restrict__ scalar_bias, const float* __restrict__ vector_weight, const float* __restrict__ gate_weight, const float* __restrict__ gate_bias, float* __restrict__ scalar_out, float* __restrict__ vector_out, int B, int N ) { const int total_items = B * N; const int threads_per_item = TILE_M; const int items_per_block = blockDim.x / threads_per_item; int item_idx = blockIdx.x * items_per_block + threadIdx.x / threads_per_item; int lane = threadIdx.x % threads_per_item; if (item_idx >= total_items) return; const float* s_in = scalar_feat + item_idx * SI; const float* v_in = vector_feat + item_idx * VI * 3; float* s_out_ptr = scalar_out + item_idx * SO; float* v_out_ptr = vector_out + item_idx * VO * 3; extern __shared__ float shared_mem[]; float* combined_shared = shared_mem; float* v_shared = shared_mem + items_per_block * (SI + VI); int item_local = threadIdx.x / threads_per_item; int combined_offset = item_local * (SI + VI); int v_offset = item_local * VI * 3; // Load scalar inputs #pragma unroll for (int i = lane; i < SI; i += threads_per_item) { combined_shared[combined_offset + i] = s_in[i]; } // Load vector inputs #pragma unroll for (int i = lane; i < VI * 3; i += threads_per_item) { v_shared[v_offset + i] = v_in[i]; } __syncthreads(); // Compute norm float v0 = v_shared[v_offset + lane * 3 + 0]; float v1 = v_shared[v_offset + lane * 3 + 1]; float v2 = v_shared[v_offset + lane * 3 + 2]; float norm = sqrtf(v0 * v0 + v1 * v1 + v2 * v2); combined_shared[combined_offset + SI + lane] = norm; __syncthreads(); // Scalar linear: outputs lane, lane+16, lane+32, lane+48 #pragma unroll for (int out_idx = 0; out_idx < SO / threads_per_item; out_idx++) { int j = lane + out_idx * threads_per_item; float sum = scalar_bias[j]; #pragma unroll for (int i = 0; i < SI + VI; i++) { sum += combined_shared[combined_offset + i] * scalar_weight[j * (SI + VI) + i]; } s_out_ptr[j] = sum; } // Gate linear float gate_sum = gate_bias[lane]; #pragma unroll for (int i = 0; i < SI + VI; i++) { gate_sum += combined_shared[combined_offset + i] * gate_weight[lane * (SI + VI) + i]; } float gate_val = 1.0f / (1.0f + expf(-gate_sum)); // Vector transform #pragma unroll for (int j = 0; j < 3; j++) { float sum = 0.0f; #pragma unroll for (int k = 0; k < VI; k++) { sum += v_shared[v_offset + k * 3 + j] * vector_weight[lane * VI + k]; } v_out_ptr[lane * 3 + j] = gate_val * sum; } } std::vector<torch::Tensor> se3_layer_forward( torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_weight, torch::Tensor scalar_bias, torch::Tensor vector_weight, torch::Tensor gate_weight, torch::Tensor gate_bias ) { int B = scalar_feat.size(0); int N = scalar_feat.size(1); int SI = scalar_feat.size(2); int VI = vector_feat.size(2); int SO = scalar_weight.size(0); int VO = vector_weight.size(0); auto scalar_out = torch::empty({B, N, SO}, scalar_feat.options()); auto vector_out = torch::empty({B, N, VO, 3}, vector_feat.options()); int total_items = B * N; const int threads_per_item = TILE_M; int items_per_block = 128 / threads_per_item; int num_blocks = (total_items + items_per_block - 1) / items_per_block; int shared_mem_size = items_per_block * ((SI + VI) + VI * 3) * sizeof(float); if (SI == 64 && SO == 64 && VI == 16 && VO == 16) { se3_layer_kernel<64, 64, 16, 16><<<num_blocks, 128, shared_mem_size>>>( scalar_feat.data_ptr<float>(), vector_feat.data_ptr<float>(), scalar_weight.data_ptr<float>(), scalar_bias.data_ptr<float>(), vector_weight.data_ptr<float>(), gate_weight.data_ptr<float>(), gate_bias.data_ptr<float>(), scalar_out.data_ptr<float>(), vector_out.data_ptr<float>(), B, N ); } else { // Fallback to generic version is not implemented in this kernel // For other sizes, we still need to handle them // Let me add a non-templated version or another template instantiation } return {scalar_out, vector_out}; } """ # ... ``` Wait, the issue is that `load_inline` function binding calls our C++ function from Python. If the function takes a `std::vector<torch::Tensor>`, how does pybind handle that? Usually, pybind11 can bind functions returning `std::vector<at::Tensor>` or `std::vector<torch::Tensor>`. But `torch::Tensor` might need proper casting. Let me use `at::Tensor` instead, or just return `torch::TensorList` or something. Actually, in custom CUDA extensions, returning a `std::vector<torch::Tensor>` is fine with pybind11. Let me verify. Alternatively, I can return a single tensor by concatenating, but that would be wasteful. The tuple return is better. Actually, the simplest approach is to have two separate kernels, or modify the C++ to return two tensors by reference. But with auto-generated pybind, we just return a std::vector. Let me check: in PyTorch C++ extensions, returning `std::vector<at::Tensor>` from a function exposed to Python creates a Python tuple of Tensors. This is standard and works well. Let me also think about whether I need a generic fallback. The test inputs are fixed at 64, 64, 16, 16, but the constructor should accept arbitrary arguments. If someone calls with different sizes, the kernel should still work. To make it robust, I can either: 1. Add runtime-parameter versions of the kernel 2. Use a non-templated kernel with runtime parameters (slower but works for all sizes) 3. Instantiating multiple template versions for common sizes For simplicity, I'll write a non-templated kernel with runtime parameters. Even without compile-time loop bounds, a well-written kernel should still be much faster than the reference due to fusion. Actually, wait. The `#pragma unroll` requires compile-time loop bounds. Without them, the compiler won't unroll, and there will be loop overhead. For loops of size 80, 16, 64, etc., loop overhead is negligible. The main performance gains come from fusing the operations, not from loop unrolling. Let me write a non-templated kernel. I'll keep the `#pragma unroll` where the bounds are compile-time constant (like `SO / threads_per_item` where `SO = 64` and `threads_per_item = 16`, giving 4, which is a compile-time constant if we write `4` explicitly). Hmm, but in a non-templated kernel, `SO` and `threads_per_item` are not compile-time constants. However, `threads_per_item` is `#define TILE_M 16`, so it is a compile-time constant. And `SO / TILE_M` is only a compile-time constant if `SO` is too. If `SO` is a function parameter, it's not. So in a non-templated kernel, `SO / TILE_M` is not known at compile time. The loop `for (int out_idx = 0; out_idx < SO / TILE_M; out_idx++)` won't be unrolled. But we can replace it with `for (int out_idx = 0; out_idx < 4; out_idx++)` if we know SO=64 and TILE_M=16. But in a generic kernel, we don't know SO. I think the best solution is to write a generic kernel that handles all sizes (good enough due to fusion) but also write a specialized kernel for the specific sizes used in the test. Then dispatch in the host function. Actually, I can simplify: for this problem, we know the test sizes. I'll write a specialized kernel for 64,64,16,16 and also a generic one for safety. Let me think about whether to use templates for the specialized kernel. Templates are easy: ```cuda template <int SI, int SO, int VI, int VO> __global__ void se3_layer_kernel_templated(...) { ... } // Generic non-templated __global__ void se3_layer_kernel_generic(...) { ... } ``` The templated kernel can use compile-time constants for all loop bounds, allowing full unrolling. For 64,64,16,16: - SI + VI = 80 - SO / TILE_M = 4 - VI * 3 = 48 - VO = 16 All loops can be fully unrolled with `#pragma unroll`: - Loading scalars: loop 64/16 = 4 iterations - Loading vectors: loop 48/16 = 3 iterations - Scalar linear: 4 outputs * 80 inputs = 4*80 = 320 iterations - Gate: 80 iterations - Vector transform: 3 * 16 = 48 iterations Total unrolled instructions: ~450 FMAs + loads per thread. This is very manageable. Let me write the templated kernel and a generic fallback. ```cuda template <int SI, int SO, int VI, int VO> __global__ void se3_layer_kernel_tmpl(...) { ... } __global__ void se3_layer_kernel_generic(...) { ... } ``` The generic kernel will be similar but without compile-time bounds. I'll dispatch in the host Python wrapper. Actually, a much simpler approach: I can use C++ `if constexpr` or a switch statement in the C++ host function: ```cpp std::vector<torch::Tensor> se3_layer_forward(...) { // ... if (scalar_in == 64 && scalar_out == 64 && vector_in == 16 && vector_out == 16) { se3_layer_kernel_tmpl<64, 64, 16, 16><<<...>>>(...); } else { se3_layer_kernel_generic<<<...>>>(...); } // ... } ``` This is clean. Let me now write the complete code. One more optimization: `__ldg()` for read-only global memory. This is a hint that the memory is read-only and can be cached in the read-only data cache. We can use it for weights and inputs. ```cuda float w = __ldg(&scalar_weight[idx]); ``` But `__ldg` only works for global memory. Using it for weights and inputs might help. Actually, on modern GPUs (Volta and later), there's no separate read-only cache; `__ldg` just uses the normal L1/L2 cache with a different policy. But it can still help with compiler optimizations. I'll keep it simple and use regular loads. Another thought: the `sqrtf` function. We can use `std::sqrt` or `sqrtf`. On CUDA, `sqrtf` is the fast single-precision sqrt. We should use `sqrtf` for fp32. Actually, for numerical accuracy, using `sqrtf` should match PyTorch's behavior. PyTorch uses `std::sqrt` in C++ which for float calls `sqrtf`. So it should match. For the sigmoid, PyTorch uses the standard formula `1 / (1 + exp(-x))`. But we need to be careful about numerical stability for large negative or positive x. PyTorch's sigmoid implementation uses: ```cpp float z = expf(-fabsf(x)); return (x >= 0) ? 1 / (1 + z) : z / (1 + z); ``` Or something similar to avoid overflow in `exp`. For our purposes, standard sigmoid should be fine since inputs are typically not extreme. Actually, let me check if PyTorch's implementation matches exactly. The tolerance is `atol=rtol=1e-4` for fp32. Even if our sigmoid has slightly different numerical behavior, it should still be within 1e-4. But to be safe, I'll use the numerically stable version: ```cuda float sigmoid(float x) { return (x >= 0.0f) ? 1.0f / (1.0f + expf(-x)) : expf(x) / (1.0f + expf(x)); } ``` Wait, if x is very negative, `expf(x)` is small, `expf(x) / (1 + expf(x))` ≈ `expf(x)`. If x is very positive, `expf(-x)` is small, `1 / (1 + expf(-x))` ≈ 1. Both avoid overflow. For x = 0: first formula gives 1/(1+1) = 0.5. Second gives 1/(1+1) = 0.5. Let me unify: ```cuda float sigmoid(float x) { if (x >= 0.0f) { float z = expf(-x); return 1.0f / (1.0f + z); } else { float z = expf(x); return z / (1.0f + z); } } ``` This is numerically stable. But it introduces a branch. On a warp with mixed signs, there would be divergence. However, for typical neural network inputs, signs are usually mixed. PyTorch's sigmoid might use a custom formula that's branchless. Let me check if I can use `1.0f / (1.0f + expf(-x))` and still be within tolerance. I'll just use the standard formula and see. The inputs to sigmoid come from a linear layer on normal-distributed data, so they'll be moderate. `expf(10)` is fine. `expf(-100)` is 0. So for typical values it should work. Actually, let me use the safer version to be sure. One more thing: should we use `Tanh` instead? No, reference uses sigmoid. OK, let me finalize the code. I'll write the non-templated generic version first as the main implementation, and add the templated version as an optimization path. Actually, wait. Maybe I should just write the generic version. The difference in performance between generic and templated might not be that significant given the small sizes. But the full unrolling could help the compiler schedule instructions better. Let me do both. Wait, I just realized something about the non-templated shared memory code. If `SI + VI` and `VI * 3` are not compile-time constants, then the shared memory indexing using runtime variables might be slightly slower. But shared memory offsets are just address calculations, which are cheap. Let me also think about whether to process items per block differently. With dynamic shared memory, the size depends on `items_per_block`. If we know the block size is 128 and items_per_block = 8, then shared memory is fixed per block. But I compute it as `items_per_block * ((SI + VI) + VI * 3) * sizeof(float)`. This is fine. Now, let me also check if we're using too much shared memory. The max shared memory per block is 227KB on H100. For our kernel with 8 items per block and 80+48=128 floats per item: 1024 floats = 4KB. Very small. One more optimization: since `items_per_block` is always 8 (because blockDim.x is 128 and threads_per_item is 16), I can hardcode it to simplify the kernel. But the host function computes it as `128 / threads_per_item`. Since threads_per_item is #define'd to 16, this is 8. So I can use compile-time constant 8 in the kernel if I want. Actually, I'll keep `items_per_block` as `blockDim.x / threads_per_item` in the kernel for clarity, and since `blockDim.x` is typically known at compile time for the templated version (but actually `blockDim.x` is a runtime special register, so it's not a compile-time constant). Hmm, `blockDim.x` is a built-in variable. The compiler knows it's constant within a kernel for a given launch configuration? No, `blockDim.x` is passed as part of the kernel launch and is available as a constant in the PTX for that launch. But the CUDA compiler generates code that treats it as semi-constant... actually no, it's a special register. In any case, using `blockDim.x / TILE_M` in the kernel means the compiler doesn't know the value at compile time. But since shared memory uses dynamic allocation via the kernel launch parameter, the compiler treats `shared_mem` as dynamically sized anyway. Let me not overthink this and just write the code. ### Final Code Plan I will write a kernel that is a hybrid: templated for compile-time sizes but with a generic fallback. Actually, to keep it simple and reduce the risk of compilation errors, I'll write a single non-templated kernel first, test it, and then add templates if needed. Here's my first draft of the complete code: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline cuda_source = ''' #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> #define TILE_M 16 __global__ void se3_layer_kernel( const float* __restrict__ scalar_feat, const float* __restrict__ vector_feat, const float* __restrict__ scalar_weight, const float* __restrict__ scalar_bias, const float* __restrict__ vector_weight, const float* __restrict__ gate_weight, const float* __restrict__ gate_bias, float* __restrict__ scalar_out, float* __restrict__ vector_out, int B, int N, int SI, int SO, int VI, int VO ) { const int total_items = B * N; const int threads_per_item = TILE_M; const int items_per_block = blockDim.x / threads_per_item; int item_idx = blockIdx.x * items_per_block + threadIdx.x / threads_per_item; int lane = threadIdx.x % threads_per_item; if (item_idx >= total_items) return; const float* s_in = scalar_feat + item_idx * SI; const float* v_in = vector_feat + item_idx * VI * 3; float* s_out_ptr = scalar_out + item_idx * SO; float* v_out_ptr = vector_out + item_idx * VO * 3; extern __shared__ float shared_mem[]; float* combined_shared = shared_mem; float* v_shared = shared_mem + items_per_block * (SI + VI); int item_local = threadIdx.x / threads_per_item; int combined_offset = item_local * (SI + VI); int v_offset = item_local * VI * 3; // Load scalar inputs cooperatively for (int i = lane; i < SI; i += threads_per_item) { combined_shared[combined_offset + i] = s_in[i]; } // Load vector inputs cooperatively for (int i = lane; i < VI * 3; i += threads_per_item) { v_shared[v_offset + i] = v_in[i]; } __syncthreads(); // Compute vector norm for this thread's vector float v0 = v_shared[v_offset + lane * 3 + 0]; float v1 = v_shared[v_offset + lane * 3 + 1]; float v2 = v_shared[v_offset + lane * 3 + 2]; float norm = sqrtf(v0 * v0 + v1 * v1 + v2 * v2); combined_shared[combined_offset + SI + lane] = norm; __syncthreads(); // Scalar linear: each thread handles 4 outputs int num_scalar_outs_per_thread = SO / threads_per_item; for (int out_idx = 0; out_idx < num_scalar_outs_per_thread; out_idx++) { int j = lane + out_idx * threads_per_item; float sum = scalar_bias[j]; for (int i = 0; i < SI + VI; i++) { sum += combined_shared[combined_offset + i] * scalar_weight[j * (SI + VI) + i]; } s_out_ptr[j] = sum; } // Gate linear: each thread handles 1 output float gate_sum = gate_bias[lane]; for (int i = 0; i < SI + VI; i++) { gate_sum += combined_shared[combined_offset + i] * gate_weight[lane * (SI + VI) + i]; } float gate_val = 1.0f / (1.0f + expf(-gate_sum)); // Vector transform: each thread handles 1 row (3 elements) for (int j = 0; j < 3; j++) { float sum = 0.0f; for (int k = 0; k < VI; k++) { sum += v_shared[v_offset + k * 3 + j] * vector_weight[lane * VI + k]; } v_out_ptr[lane * 3 + j] = gate_val * sum; } } std::vector<at::Tensor> se3_layer_forward( at::Tensor scalar_feat, at::Tensor vector_feat, at::Tensor scalar_weight, at::Tensor scalar_bias, at::Tensor vector_weight, at::Tensor gate_weight, at::Tensor gate_bias ) { int B = scalar_feat.size(0); int N = scalar_feat.size(1); int SI = scalar_feat.size(2); int VI = vector_feat.size(2); int SO = scalar_weight.size(0); int VO = vector_weight.size(0); auto scalar_out = torch::empty({B, N, SO}, scalar_feat.options()); auto vector_out = torch::empty({B, N, VO, 3}, vector_feat.options()); int total_items = B * N; const int threads_per_item = TILE_M; int items_per_block = 128 / threads_per_item; int num_blocks = (total_items + items_per_block - 1) / items_per_block; int shared_mem_size = items_per_block * ((SI + VI) + VI * 3) * sizeof(float); se3_layer_kernel<<<num_blocks, 128, shared_mem_size>>>( scalar_feat.data_ptr<float>(), vector_feat.data_ptr<float>(), scalar_weight.data_ptr<float>(), scalar_bias.data_ptr<float>(), vector_weight.data_ptr<float>(), gate_weight.data_ptr<float>(), gate_bias.data_ptr<float>(), scalar_out.data_ptr<float>(), vector_out.data_ptr<float>(), B, N, SI, SO, VI, VO ); return {scalar_out, vector_out}; } ''' # Compile se3_cuda = load_inline( name="se3_cuda", cpp_sources="", cuda_sources=cuda_source, functions=["se3_layer_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False, ) class ModelNew(nn.Module): def __init__(self, scalar_in, scalar_out, vector_in, vector_out): super().__init__() self.scalar_in = scalar_in self.scalar_out = scalar_out self.vector_in = vector_in self.vector_out = vector_out self.scalar_linear = nn.Linear(scalar_in + vector_in, scalar_out) self.vector_weight = nn.Parameter(torch.randn(vector_out, vector_in)) nn.init.xavier_uniform_(self.vector_weight) self.vector_gate = nn.Linear(scalar_in + vector_in, vector_out) def forward(self, scalar_feat, vector_feat): # Extract weights from PyTorch layers scalar_weight = self.scalar_linear.weight # (scalar_out, scalar_in + vector_in) scalar_bias = self.scalar_linear.bias # (scalar_out,) gate_weight = self.vector_gate.weight # (vector_out, scalar_in + vector_in) gate_bias = self.vector_gate.bias # (vector_out,) return se3_cuda.se3_layer_forward( scalar_feat, vector_feat, scalar_weight, scalar_bias, self.vector_weight, gate_weight, gate_bias ) ``` Wait, there's a CUDA runtime error potential. `load_inline` compiles at module load time. The kernel assumes contiguous float tensors. We should ensure the inputs are contiguous float tensors. Also, `.data_ptr<float>()` requires the tensor to be of float type. The inputs should be float32. One more issue: the `se3_layer_kernel` only works if: - `SO % TILE_M == 0` (each thread gets an equal number of scalar outputs) - `SI % TILE_M == 0` or the load loop works correctly - `VI == TILE_M` (exactly 16 vectors, each thread handles exactly one) For the generic case where VI might not equal 16, the kernel as written assumes lane < VI for norm computation, gate, and vector transform. If VO < 16, we'd have issues too. But for the test, VI=16 and VO=16, and SI=64, SO=64. So SO/TILE_M = 4. It works. To make it robust, for non-power-of-2 sizes, I should add bounds checking. But for the given problem, the sizes match perfectly. Actually, I just need to make sure that `vector_in * 3` is evenly divisible by `threads_per_item` for the cooperative load to work without gaps. `16 * 3 = 48`, and `48 / 16 = 3`. So each thread loads exactly 3 floats. Good. For scalar_in = 64, each thread loads `64 / 16 = 4` floats. Good. If sizes were different (e.g., scalar_in = 65), the load loop `for (int i = lane; i < SI; i += threads_per_item)` handles it correctly. But then `SI + VI` might not divide evenly for the dot products, and `SO / threads_per_item` might not be integer. For the generic case, I could add a bounds check in the kernel. For example: ```cuda if (lane >= VI) return; // or similar ``` But then within a warp, some threads would exit early, causing divergence. However, for valid sizes where lane < VI and lane < VO, all threads stay active. Since the problem specifies valid sizes where VI=VO=16 and SI=SO=64, and our threads_per_item is 16, all 16 threads in each group will be active for all operations. Let me also consider the case where SO is not a multiple of 16. We'd need to handle the remainder. But SO=64 is a multiple of 16, so it's fine. To be safe, I'll modify the kernel to handle arbitrary sizes gracefully. But I might keep it simple for the common case. One more issue: I need to make sure `scalar_feat` and `vector_feat` are contiguous. Non-contiguous tensors from slicing could cause issues. I'll call `.contiguous()` in the Python wrapper. Also, I should check if the tensors are on CUDA. They should be, given the problem setup. Let me now compile and test this kernel. Actually, about the `std::vector<at::Tensor>` return type: I need to use `at::Tensor` consistently. In the C++ code, `torch::Tensor` and `at::Tensor` are the same type in the PyTorch C++ API. But I should use one consistently. Let me use `at::Tensor` for the function signature since we're in the ATen library context. Wait, actually looking at PyTorch C++ extension examples, people typically use `torch::Tensor` interchangeably with `at::Tensor`. Both work. I'll use `torch::Tensor` for consistency with `torch::empty`. Also, I should consider whether the weights are stored as row-major and whether my indexing is correct. PyTorch stores tensors in row-major order by default. A weight of shape (64, 80) is stored such that `weight[i, j]` is at `data_ptr[i * 80 + j]`. So `scalar_weight[j * (SI + VI) + i]` accesses column `i` of row `j`. This is correct for our dot product where output `j` needs all inputs `i`. For `vector_weight[lane * VI + k]`, vector_weight has shape (VO, VI) = (16, 16). `lane * VI + k` accesses column `k` of row `lane`. This is correct for the einsum `C[o,i] * V[i,d]` where `o = lane`. OK, indexing looks correct. One more thing: the CUDA kernel might be compiling with warnings about unused variables or potential issues. Let me make sure all variables are used correctly. I should also add `const` to appropriate pointers to help the optimizer. Now, let me prepare the code and compile it. But first, I realize that `load_inline` with empty `cpp_sources` should be fine; all code goes in `cuda_sources`. Wait, one important note: the C++ wrapper function must have `extern "C"` or not? Actually no, C++ name mangling is fine for pybind11. `load_inline` generates the pybind11 bindings automatically. But I should make sure the function signature uses types that pybind11 can bind. `std::vector<torch::Tensor>` should work fine. Let me write the complete code. Actually, I realize I might need to handle the fact that `scalar_weight`, `scalar_bias`, etc. come from PyTorch modules. When we pass them to the CUDA kernel, they need to be contiguous and on the same device as the input tensors. `nn.Linear.weight` is already contiguous by default. `nn.Linear.bias` is contiguous. `nn.Parameter` is contiguous. So we don't need `.contiguous()`. But to be safe, we can add it - though for parameters on CUDA, `.contiguous()` is a no-op if already contiguous. Now let me write the code and compile. Wait, I need to be careful about loading the weights. In a real `nn.Module`, `self.scalar_linear.weight` has `requires_grad=True`. But we're only doing forward pass, so that's fine. One more question: are the weights in fp32? By default, yes. OK, let me compile. Hmm, actually I realize there might be an issue with `scalar_weight.data_ptr<float>()` if the tensor isn't contiguous. Let me not worry about this since the weights should be contiguous. Also, `vector_feat` shape is `(B, N, vector_in, 3)` = `(8, 128, 16, 3)` = 49152 elements. This is contiguous. And `scalar_feat` is `(8, 128, 64)` = 65536 elements. Contiguous. OK here's the code. I'll include `#include <vector>` for std::vector but it's probably already included by torch/extension.h. Actually, torch/extension.h includes everything necessary. One potential issue: dynamic shared memory allocation in CUDA must be specified in the kernel launch: `<<<grid, block, shared_mem_size>>>`. I do this correctly. Also, `extern __shared__ float shared_mem[];` requires the shared memory to be sized correctly. In my calculation: `int shared_mem_size = items_per_block * ((SI + VI) + VI * 3) * sizeof(float);` For SI=64, VI=16, items_per_block=8: shared_mem_size = 8 * (80 + 48) * 4 = 8 * 128 * 4 = 4096 bytes. Correct. Now, the `v_shared` pointer in shared memory is at offset `items_per_block * (SI + VI)`. For 8 items * 80 = 640 floats = 2560 bytes. So v_shared starts at byte 2560. It has size `items_per_block * VI * 3 = 8 * 48 = 384` floats = 1536 bytes. Total: 2560 + 1536 = 4096 bytes. Great. Let me make sure there's no shared memory bank conflict in the loop: `for (int i = lane; i < SI; i += threads_per_item)` With `lane` in 0..15 and `threads_per_item=16`, and SI=64: Iteration 0: threads 0..15 write to shared[0..15] Iteration 1: threads 0..15 write to shared[16..31] ... These are all perfectly contiguous within each iteration, no bank conflicts. For `v_shared[v_offset + lane * 3 + j]`: Thread lane accesses `v_offset + lane * 3 + 0`, `v_offset + lane * 3 + 1`, `v_offset + lane * 3 + 2`. Addresses are `lane*12`, `lane*12+4`, `lane*12+8` bytes relative to v_offset. Bank = (lane*12/4) % 32 = lane*3 % 32. For lane=0..15: banks = 0,3,6,9,12,15,18,21,24,27,30,1,4,7,10,13. All distinct! No bank conflicts. For `combined_shared[combined_offset + i]` where all threads access the same `i` in a loop iteration: This is a broadcast. On modern GPUs, shared memory broadcast has no bank conflicts (or minimal). Great! For weights in global memory, we have scattered access as discussed, but weights are small and cached. Let me write and compile the code.

The issue seems to be that `load_inline` expects the function declaration in the cpp_sources, not the cuda_sources. Let me fix the code structure and provide a proper declaration. Also the scale is off: TILE_M=16 mapping for thread assignment on VI is wrong. VO and VI are both 16 so lane goes from 0 to 15, that's fine. But later I use `lane` for scalar out too, `SO=64`, need separate approach. Better approach: I'll just simplify, avoid the shared mem tile complexities and just use element-level parallelism. The kernel is memory bandwidth bound anyway. Let me write a cleaner kernel.

We have a kernel that compiles. But there are some issues: 1. We run a scalar kernel AND a vector kernel, both recomputing norms from vector data redundantly 2. The vector kernel writes vector_out but receives unnecessary arguments (scalar_weight, scalar_bias, scalar_out) 3. Need to check correctness Let me think more carefully about the kernel design. This problem has: - scalar_feat: (B, N, scalar_in) - vector_feat: (B, N, vector_in, 3) Operations: 1. vec_norm = vector_feat.norm(dim=-1) # (B, N, vector_in) 2. combined = cat([scalar_feat, vec_norm], -1) 3. s_out = linear(combined) -> scalar_out 4. v_out = einsum("oi,bnid->bnod", vector_weight, vector_feat) # (B, N, vector_out, 3) 5. gate = sigmoid(linear_gate(combined)).unsqueeze(-1) # (B, N, vector_out, 1) 6. v_out = gate * v_out For correctness first, we should just validate this. Let's run correctness check.

The first version compiles and passes correctness. Now let me think about how to optimize this. The current implementation has two separate kernels, each computing norms redundantly. The problem is memory bandwidth bound - we have to read the same inputs multiple times. Let's think about fusing everything into a single kernel that can operate efficiently. For H100, I should think about: 1. Fusing the scalar and vector computations to avoid redundant norm computations and memory bandwidth 2. Using shared memory for weights to improve reuse 3. Better thread mapping - perhaps using block-level mapping where each block handles multiple (B,N) elements The dimensions: - B=8, N=128, so BN=1024 - scalar_in=64, scalar_out=64 - vector_in=16, vector_out=16 Total work: - scalar: 1024 * 64 = 65536 outputs - vector: 1024 * 16 * 3 = 49152 outputs For each (B,N) pair, we need to: 1. Read SI=64 scalar features 2. Read VI*3=48 vector features 3. Compute VI=16 norms 4. For each of SO=64 outputs: linear combination over 80 features (64+16) 5. For each of VO=16 outputs: compute gate from 80 features, compute 3 x linear combo of VI=16 vectors, multiply by gate The key insight is that we can do all this per (B,N) element and keep everything in registers/shared memory. Since BN=1024, and each element has relatively small dimension, we should map (B,N) to blocks and let threads within a block compute the outputs. Let me think about a better design: - Each block handles one (B,N) element - Threads within a block compute the outputs in parallel - Use shared memory for the combined input features [SI + VI] But with 1024 elements and 64 SMs on H100, we want more parallelism. Let's have each handle handle multiple elements or use a grid-stride loop. Actually, a simpler approach: each block handles a single (B,N) element. - Block size: 128 threads - Grid: 1024 blocks - Each thread computes some fraction of SO and VO outputs For scalar outputs: SO=64, so each thread could compute 1 output with 64 threads, but we have 128 threads. For vector outputs: VO=16, 3 components each = 48 outputs. With 128 threads: - 64 threads for 64 scalar outputs (1 each) - 64 threads for 16 vector outputs (16 * 3 = 48 outputs, so each thread can do some) Wait, but we need to synchronize for shared memory. Let me think more carefully. Actually, a simpler fused kernel: - Each thread handles one (b, n, o) for vector outputs and one (b, n, o') for scalar outputs? That's messy. Alternative: one thread handles one (b, n) element completely. That's 65536 threads total (for scalar) or 49152 (for vector) or just 1024 if each handles everything. But 1024 threads is too few. Better: thread block handles one (b, n) element. - 128 threads - First 64 threads each compute one scalar output - All 128 threads participate in loading data to shared memory - Then 48 threads (or 64) each compute one vector component output Wait, for vector there are VO*3 = 48 outputs. With 128 threads, we could use 48 threads for vector and 64 for scalar, with some overlap. Actually, let me try a simpler approach: fusion of scalar and vector into a single kernel with each thread computing one output element, either scalar or vector. But the norms would be computed twice (once for scalar kernel, once for vector kernel). The current version does two kernels which each compute norms. A single fused kernel per (B,N) element would compute norms once and use them for both scalar and vector, saving memory bandwidth and compute. Let me design a block-level fused kernel: - Block size: 128 threads - Each block processes one (b, n) element - Shared memory: combined[SI+VI] for input features Algorithm: 1. All threads load scalar features into shared memory 2. All threads load vectors, compute norms, store norms in shared memory 3. __syncthreads() 4. Threads 0..SO-1 compute scalar linear outputs 5. Threads 0..VO*3-1 compute vector outputs (gate + linear) But we need to handle both scalar and vector outputs with the same threads. Actually, we can split: - 64 threads for scalar outputs (each thread computes 1 scalar output, for SO=64) - 64 threads for vector outputs (each thread computes some vector outputs) For vector: VO=16 outputs, each with 3 components. That's 48 outputs total. With 64 threads, some threads are idle. Actually wait, if SO=64 and VO=16, then: - 64 threads handle 64 scalar outputs (1 per thread) - 64 threads handle 48 vector components. But actually, each VO output needs its own gate value, and then 3 components. For vector part: - Thread tid (0 to 47 or distribute among 64 threads) - Actually, maybe better: 16 threads each handle one VO output (computing gate + 3 components) Let me redesign: - Warp-level optimization. A warp is 32 threads. - 2 warps (64 threads) handle scalar outputs - 2 warps (64 threads) handle vector outputs, but VO=16 and we have 64 threads, so each thread can handle part of fewer outputs. Actually, let's think about it differently: - Block size = 128 - First 64 threads (0-63) each compute 1 scalar output - Threads 64-79 (16 threads) each compute 1 vector output (all 3 components + gate) - Threads 80-127 are idle This seems suboptimal. Instead, let me just let the scalar and vector parts be done by the same threads at different times, or use a block size of 64 and launch more blocks by having each block handle one (b, n). With 1024 blocks and 64 threads each, we get 65536 threads total, which is fine. Wait, actually even better: we can launch 1024 blocks with 128 threads, and each block handles one (b,n). But the problem is that we need to compute both scalar and vector outputs. Let me just have all 128 threads compute everything: Phase 1: load data to shared memory (all 128 threads participate) Phase 2: compute scalar outputs (first 64 threads), compute vector outputs (some threads) Phase 3: write outputs But this gets complicated with thread divergence. Alternatively, I can keep it simple: just use 1D thread mapping where each thread handles one (b, n, o) for scalar and one (b, n, o/3, o%3) for vector. But we'd need to launch two separate grids within one kernel... no, CUDA doesn't allow that easily. Actually, the cleanest approach for performance: launch separate kernels but avoid redundant computation. But separate kernels mean we read the same inputs twice, unless we cache. Hmm, let me think about memory bandwidth: - Reading scalar: 1024 * 64 = 65536 floats = 256 KB - Reading vector: 1024 * 48 = 49152 floats = 192 KB - Total input: 448 KB - For two kernels: ~896 KB read - Writing scalar: 1024 * 64 = 256 KB - Writing vector: 1024 * 48 = 192 KB - Total output: 448 KB - Total memory traffic: 1344 KB For a single fused kernel: - Read inputs once: 448 KB - Write outputs: 448 KB - Total: 896 KB - That's a 1.5x reduction in memory bandwidth. Given the small problem size, this matters a lot. Let me design a proper fused kernel. Design: each block handles multiple (b, n) elements, or each block handles one (b, n) element. With block per (b,n): - 1024 blocks - Block size: 128 or 256 threads - Each block loads SI scalars + VI*3 vectors, computes VI norms - Stores combined [SI+VI] in shared memory - Computes all SO scalar outputs and all VO*3 vector outputs For a block of 128 threads handling one (b,n): - Load phase: 128 threads load 64 scalars (each loads ~0.5 floats) and 48 vectors (each loads ~0.375 floats). Not enough work to distribute, but that's fine. - Sync - Compute scalar: 64 threads, each does 80 FMAs. That's 5120 FMAs. With 128 threads, first 64 are busy, rest idle. - Compute vector: for each of 16 outputs, a thread or group of threads does 80 operations for gate + 16*3 FMAs for vectors. Total per output: ~128 ops. Total vector: 2048 ops. - Total compute per (b,n): ~7168 FMAs + exp + norms With H100 FP32: 67 TFLOPS = 67,000 GFLOPS. Each (b,n) has ~7K FLOPs. 1024 elements = ~7M FLOPs. At peak compute, this is ~0.1 microseconds. But we're definitely bandwidth bound for this small problem. Wait, actually for fused kernel, let's compute more precisely: - Norms: 16 * (3 mul + 2 add + sqrt) ~ 16 * 6 = 96 ops - Scalar: 64 * (80 FMA) = 5120 FMAs - Gate: 16 * (80 FMA + exp + div + add) = ~1280 + 16*10 = ~1440 ops - Vector linear: 16*3 * 16 FMA = 768 FMAs - Total: ~5120*2 + 768*2 + 1440 + 96 = ~13312 ops - 1024 elements: ~13.6M ops At 67 TFLOPS, compute time is ~0.2 microseconds. Memory bandwidth for fused: 896 KB. At 3.35 TB/s, memory time is ~0.27 microseconds. So we're roughly balanced. For two kernels: 1344 KB, memory time ~0.4 microseconds. Fused saves ~33% memory time. But the key issue with my current implementation isn't just the number of kernels. It's also that I'm doing element-wise launches with poor occupancy and memory access patterns. Actually, let me check if the current version is even faster than PyTorch. It probably is, given that PyTorch has the overhead of multiple separate ops. But I can likely do much better. Wait, I should also consider using a warp-level approach where a warp handles one (b,n). With 32 threads per warp, and 1024 elements, we need 32 warps. A block with 4 warps = 128 threads can handle 4 elements. So 256 blocks total. That's actually a decent occupancy. Let me design a warp-level kernel: - Each warp (32 threads) handles one (b, n) element - Block size: 128 threads (4 warps) - Elements per block: 4 - Total blocks: ceil(1024/4) = 256 Within a warp: - All 32 threads cooperate to load data (via shared memory or direct loads) - Actually, each thread can load some scalars and vectors directly from global memory - But we need norms and combined features that all threads use For scalar output: 64 outputs / 32 threads = 2 outputs per thread - Each thread computes 2 scalar linear combinations over 80 features For vector output: 16 outputs * 3 components = 48. 48 / 32 = 1.5 components per thread. - But actually each vector output has a shared gate value. - Maybe better: each thread handles one VO output (with gate + 3 components), but 16 outputs < 32 threads. Hmm, this gets complicated. Let me just try a simpler but more efficient fused block-level approach. Block-level fused kernel: - BlockDim = 128 threads - Each block handles `items_per_block` (b,n) elements - For simplicity, items_per_block = 4 (since 128 threads / 32 warps = 4 warps) Wait, but we need to store combined [SI+VI] per element in shared memory. For 4 elements: 4 * (64+16) = 320 floats. For vector inputs: 4 * 48 = 192 floats. Total shared mem: ~512 floats * 4B = 2KB. Very small! Actually, let me do items_per_block = 8. Then 128/8 = 16 threads per item. Shared mem: 8 * 80 + 8 * 48 = 1024 floats = 4KB. Still small. For each item, 16 threads: - Load 64 scalars: each thread loads 4 - Load 48 vectors: each thread loads 3 - Compute 16 norms: each thread computes ~1 - Store into shared memory - Compute 64 scalar outputs: each thread computes 4 - Compute 16 vector outputs: each thread computes 1 (gate + 3 components) This seems like a good balance! Let me write this kernel: ```cuda __global__ void fused_kernel(...) { const int items_per_block = 8; const int threads_per_item = blockDim.x / items_per_block; // 16 int item_in_block = threadIdx.x / threads_per_item; int lane = threadIdx.x % threads_per_item; int item_idx = blockIdx.x * items_per_block + item_in_block; if (item_idx >= B * N) return; // Compute offsets // Load data... // Store to shared memory... // Compute outputs... } ``` Actually, let me refine this. The issue is that with items_per_block = 8 and grid = ceil(1024/8) = 128 blocks, we only have 128 blocks. On H100 with 132 SMs and 4 warps per block, that's 512 warps total. Occupancy should be fine since each SM can handle many warps. Let me code this up. But first, I'm concerned about correctness since the current version already passes. Let me just make a more optimized version. Actually, wait - I realize the current simple version might be fine. But let me think about whether there are correctness issues with my current kernel. In the scalar kernel, I have: ```cuda float sum = scalar_bias[o]; ... scalar_out[idx] = sum; ``` In the vector kernel, I have: ```cuda float gate_val = gate_bias[o]; ... for (int j = 0; j < 3; ++j) { float vsum = 0.0f; ... vector_out[idx * 3 + j] = gate_val * vsum; } ``` Wait, in the original PyTorch model: ```python v_out = torch.einsum("oi,bnid->bnod", self.vector_weight, vector_feat) gate = torch.sigmoid(self.vector_gate(combined)).unsqueeze(-1) v_out = gate * v_out ``` So `v_out` shape is (B, N, VO, 3) and `gate` shape is (B, N, VO, 1). The gate is broadcast over the last dimension. In my kernel, `idx = bno * VO + o`, so `vector_out[idx * 3 + j]` corresponds to (b, n, o, j). And `gate_val` is the same for all j for a given (b,n,o). This is correct. The scalar part: ```python s_out = self.scalar_linear(combined) ``` This is just a linear layer. In my kernel, `scalar_out[idx]` where `idx = bns * SO + o` corresponds to (b, n, o). Correct. OK, the current implementation is correct. Now let me optimize it with a fused kernel. Let me write a cleaner, more optimized fused kernel: ```cuda __global__ void fused_se3_kernel( const float* scalar_feat, const float* vector_feat, const float* scalar_weight, const float* scalar_bias, const float* vector_weight, const float* gate_weight, const float* gate_bias, float* scalar_out, float* vector_out, int B, int N, int SI, int SO, int VI, int VO ) { const int items_per_block = 8; const int threads_per_item = blockDim.x / items_per_block; int item_local = threadIdx.x / threads_per_item; int tid = threadIdx.x % threads_per_item; int item_idx = blockIdx.x * items_per_block + item_local; if (item_idx >= B * N) return; extern __shared__ float sdata[]; float* combined = sdata + item_local * (SI + VI); float* vectors = sdata + items_per_block * (SI + VI) + item_local * (VI * 3); const float* s_in = scalar_feat + item_idx * SI; const float* v_in = vector_feat + item_idx * (VI * 3); // Load scalars for (int i = tid; i < SI; i += threads_per_item) { combined[i] = s_in[i]; } // Load vectors for (int i = tid; i < VI * 3; i += threads_per_item) { vectors[i] = v_in[i]; } __syncthreads(); // Compute norms for (int i = tid; i < VI; i += threads_per_item) { float v0 = vectors[i * 3 + 0]; float v1 = vectors[i * 3 + 1]; float v2 = vectors[i * 3 + 2]; combined[SI + i] = sqrtf(v0*v0 + v1*v1 + v2*v2); } __syncthreads(); float* s_out_ptr = scalar_out + item_idx * SO; float* v_out_ptr = vector_out + item_idx * (VO * 3); // Scalar linear outputs for (int o = tid; o < SO; o += threads_per_item) { float sum = scalar_bias[o]; for (int i = 0; i < SI + VI; ++i) { sum += combined[i] * scalar_weight[o * (SI + VI) + i]; } s_out_ptr[o] = sum; } // Vector outputs: gate + linear for (int o = tid; o < VO; o += threads_per_item) { float gate = gate_bias[o]; for (int i = 0; i < SI + VI; ++i) { gate += combined[i] * gate_weight[o * (SI + VI) + i]; } gate = 1.0f / (1.0f + expf(-gate)); for (int j = 0; j < 3; ++j) { float vsum = 0.0f; for (int i = 0; i < VI; ++i) { vsum += vectors[i * 3 + j] * vector_weight[o * VI + i]; } v_out_ptr[o * 3 + j] = gate * vsum; } } } ``` Wait, I need to double check: `item_idx * SI` for scalar_feat - scalar_feat shape is (B, N, SI), and its stride should indeed be N*SI or just SI depending on whether it's contiguous. Since I've called `.contiguous()`, the memory layout is (B*N*SI) with scalar_feat[b, n, i] at index (b*N + n)*SI + i. That's `item_idx * SI + i`. Yes, correct. For vector_feat: shape (B, N, VI, 3), vector_feat[b, n, i, j] at index ((b*N + n)*VI + i)*3 + j = item_idx * (VI*3) + i*3 + j. Yes, correct. For scalar_out: same as scalar_feat. Correct. For vector_out: same layout as vector_feat. Correct. For scalar_weight: shape (SO, SI+VI). Correct. For vector_weight: shape (VO, VI). Correct. For gate_weight: shape (VO, SI+VI). Correct. OK, now one more optimization: we could keep weights in shared memory too. Since weights are small: - scalar_weight: SO * (SI+VI) = 64 * 80 = 5120 floats = 20KB - vector_weight: VO * VI = 16 * 16 = 256 floats = 1KB - gate_weight: VO * (SI+VI) = 16 * 80 = 1280 floats = 5KB Total weights: ~26KB. With 8 items per block using the data arrays, plus weights: total shared mem per block would be about 8 * 80 * 4 + 8 * 48 * 4 + 26KB = ~2.5KB + 1.5KB + 26KB = ~30KB. That's well within the 227KB limit. Loading weights to shared memory once per block could help with memory access patterns. Actually, weights are read by all blocks. But with the small problem size, they might already be in L2 cache. Hmm, on H100 the L2 cache is 50MB or more, so all weights easily fit. So loading weights to shared memory isn't critical, but it might help with coalescing. Actually, looking at the weight access pattern: - In my current kernels, weights are loaded from global memory with good coalescing because adjacent threads access adjacent elements (o * (SI+VI) + i where o differs by 1). Actually, for scalar outputs, thread with idx computes output o. So threads with adjacent idx have adjacent o. When accessing weight[o * (SI+VI) + i], adjacent threads access adjacent weights, which is coalesced. This is good. One issue: in the loop `for (int i = 0; i < SI + VI; ++i)`, all threads read the same weight element in lockstep. This is broadcast, which is actually fast on modern GPUs because it hits L1 cache. So the current access pattern is already decent. Wait, but there's another issue: the vector computation has `for (int j = 0; j < 3; ++j)` as the outer loop and `for (int i = 0; i < VI; ++i)` as the inner loop. This means we access vector_weight[o * VI + i] for i=0..15 in the inner loop. Since o is constant for a given thread, and VI=16, we access 16 floats. But adjacent threads (different o) access different parts of vector_weight. This isn't coalesced. Actually, let me reconsider. For vector_weight, shape is (VO, VI) = (16, 16). Threads o=0 to o=15 access row 0 to row 15. For a fixed i, vector_weight[o * VI + i] is strided by VI=16. Adjacent threads access elements separated by 16 floats = 64 bytes. This isn't perfectly coalesced but it's OK. For vectors[i * 3 + j], with j varying in the inner loop (or i varying), this is coalesced since adjacent i values are adjacent in memory. Let me also think about whether I should use `__ldg` (cached loads) for the weight accesses. Since weights are read many times and are small, `__ldg` could help. Actually `__ldg` is for read-only data that can use the texture/L1 cache aggressively. Let me try using `const __restrict__` which should already enable many optimizations. One more thing: I should check if the problem is actually compute-bound or memory-bound. Let me compute arithmetic intensity. Fused kernel per (b,n): - Load: 64 + 48 = 112 floats = 448 bytes - Compute: ~10000 FLOPs - Store: 64 + 48 = 112 floats = 448 bytes - AI = 10000 / 896 = 11.2 FLOP/byte At 3.35 TB/s memory and 67 TFLOPS compute, the ridge point is 67,000 / 3.35 = 20 FLOP/byte. Our AI is 11.2, so we're actually memory-bound! That means fusing is very important to reduce memory bandwidth. With fused kernel, we read 112 floats and write 112 floats = 224 floats = 896 bytes. With two kernels, we read 112 floats twice and write separately = 336 floats read + 112 floats written = 448 floats = 1792 bytes. So fused kernel halves memory bandwidth. That's a big win. Let me also consider whether I can use vectorized loads/stores. For loading scalar and vector features, I could use float4 to load 4 floats at a time. Since SI=64 and VI*3=48, both are divisible by 4, so this works. Actually, a simpler optimization: use `float4` for vector loads where possible. Hmm, for now let me focus on getting the fused kernel right. Float4 optimization can be added later. Let me write and test the fused kernel. Actually, one more thought: in my current two-kernel design, I can see a potential issue. The vector kernel doesn't use `scalar_out`, so the compiler might optimize things differently. But it also means we launch two kernels from Python, which has some overhead (~10-20 microseconds per kernel launch). For a small problem that might take ~50 microseconds total, 2 kernel launches is significant overhead. A single fused kernel would eliminate this overhead. Let me code up the fused version and test it.

The fused kernel compiles and passes correctness. Now let me optimize further. The current kernel uses: - 128 threads per block, 16 items per block - Each item is a (B,N) point - Threads within the warp process different outputs for the same item This is actually not a great warp organization. Let me think about a better approach. For H100, we want to maximize occupancy and memory bandwidth. Let me try a few things: 1. Use a warp per item rather than 8 threads per item 2. Use vectorization for loads/stores 3. Consider using row-major weight matrices more effectively Actually, let me look at the performance. With 16 items per block and 128 threads, each warp (32 threads) covers 2 items with 16 threads each. This is already decent, but maybe I can use 8 warps per block = 4 items per block with 8 warps each = 256 threads. Wait, actually each item has: - SI + VI = 80 combined features - SO = 64 scalar outputs - VO = 16 vector outputs With 8 threads per item doing 8 outputs at a time, each thread processes only 8 outputs. For SO=64, each thread does 8 outputs, which is a nice number. But for VO=16, each thread does only 2 outputs with 8 threads. Still fine. Let me try tuning the threads per item and block configuration. Let me also try the approach of a warp per item (32 threads): - Items per block = 128 threads / 32 = 4 items - Each item gets 1 warp - For scalar (SO=64), each thread does 2 scalar outputs - For vector (VO=16), each thread does 1 vector output (with 3 components) Let me also try to register the weights, or use unrolled loops with known sizes. Better yet, let me hardcode the sizes in templates. Since the sizes are fixed (64, 64, 16, 16, 128, 8), I can use template specializations that the compiler can optimize aggressively. But `load_inline` with templates can be tricky. Let me try using constexpr if possible, or just rely on the compiler to unroll with `#pragma unroll`. Actually, I can use template parameters since they're fixed at compile time in our kernel: ```cpp template <int SI, int VI, int SO, int VO> __global__ void kernel(...) ``` And then dispatch based on runtime values using if-else or a switch. But since SI=64, VI=16, SO=64, VO=16, I can just hardcode them or use templates. Let me try a more optimized version with: 1. Warp per item (32 threads) 2. Each warp cooperatively processes all outputs for its item 3. Register caching for weights 4. Better use of shared memory for the combined features Actually, with a warp per item: - Combined array: 80 floats - Vector array: 48 floats - Per block (4 items): 320 + 192 = 512 floats = 2KB shared memory - Very small, we can fit many blocks Let me write this version. Wait, with only 4 items per block and B*N=1024 items, we'd need 256 blocks. With 128 SMs, that's 2 blocks per SM, which is not great occupancy. Let me go back to more items per block. With 16 items per block and 8 threads per item, we get: - 128 threads per block - 1024/16 = 64 blocks - 64 blocks / 128 SMs = 0.5 blocks per SM initially But as threads finish, more blocks can be scheduled. The limit is occupancy. Actually, the better metric is total threads. 64 blocks * 128 threads = 8192 threads total. H100 can run thousands of threads per SM. Let me try a different organization: one thread per item, and each thread loops over all outputs independently. This means: - 1024 threads total - Threads per block = 256 - Blocks = 4 - Very few blocks That's bad. We want at least hundreds of blocks to saturate the GPU. What about: - Each warp (32 threads) processes one item - So items_per_block = warps_per_block - Let warps_per_block = 8, so threads = 256, items = 8 - Blocks = 1024/8 = 128 This gives 128 blocks of 256 threads = 32768 threads total, which is good. Within a warp (32 threads) for one item: - 32 threads load scalars (64 floats) -> 2 loads each thread - 32 threads load vectors (48 floats) -> 1-2 loads each thread - Compute norms: 16 norms, each thread can do 1 norm (with some threads idle? No) Actually, I think the original 16 items x 8 threads is fine. Let me try tuning: - Maybe 32 items per block x 4 threads per item? - Or 8 items per block x 16 threads per item? With 8 items x 16 threads: - Each block = 128 threads - 1024/8 = 128 blocks - For SO=64, each thread handles 4 scalar outputs - For VO=16, each thread handles 1 vector output (3 components) This might be better since each vector output has 3 components that share the same gate value. A single thread computing all 3 components avoids broadcasting the gate. Let me try this configuration. Actually, wait. The current performance might be fine. Let me just try to optimize memory access patterns and instruction level parallelism. One issue: scalar_weight is accessed as `scalar_weight + o * (SI + VI) + i`. Each thread in a warp accesses the same row `o` but different `i`. This is coalesced if threads access adjacent elements. In my current code, threads within a warp have `o = tid`, and they process different `o` values, so they access different rows. This means it's NOT coalesced within a warp, because thread 0 accesses row 0, thread 1 accesses row 1, etc. The addresses are spaced by (SI+VI) floats. To make weight access coalesced, I need all threads in a warp to access the same column `i` but different rows `o`. But then `i` needs to be the same across threads. Actually wait, let's think about this more carefully: - Thread 0: accesses scalar_weight[0, i] for i=0..79 - Thread 1: accesses scalar_weight[1, i] for i=0..79 If thread 0 and thread 1 are in the same warp and executing in lockstep, they're accessing: - Step 1: thread0 -> scalar_weight[0], thread1 -> scalar_weight[80], thread2 -> scalar_weight[160], ... - Step 2: thread0 -> scalar_weight[1], thread1 -> scalar_weight[81], ... This is actually stride-80, which is bad for coalescing because the cache line is 32 floats and threads are accessing positions 0, 80, 160, ... which are 320 bytes apart. A 128-byte cache line holds only 1 of these accesses. Wait, the closest thread would access 80 floats apart. 80*4 bytes = 320 bytes. A cache line is typically 128 bytes on H100. So each access from 8 threads would span 7 cache lines. That's very bad for memory bandwidth. To fix this, I should have threads in a warp process the SAME output `o` (same row) but different elements of the input (different `i`). Then at each step, they access adjacent memory. But then I need multiple warps processing different `o` values. Let me redesign: - One warp per output `o` for a given item - Multiple warps process all outputs for the same item in parallel - Items are processed by blocks For one item: - SI + VI = 80 inputs - SO = 64 outputs -> 64 warps per item if one warp per output - VO = 16 outputs -> 16 warps per item if one warp per output This would be too many warps. Let me instead have one warp handle multiple outputs. Alternative: process outputs in tiles. Each warp handles, say, 4 outputs. Then each warp processes 4 * 80 = 320 multiplications. For scalar (SO=64): each warp handles 4 outputs, so 16 warps are needed For vector (VO=16): each warp handles 4 outputs, so 4 warps are needed Hmm, but then I need to coordinate warps within a block. Let me think of a simpler approach: - One warp handles all scalar outputs for one item: 32 threads, each thread does 2 outputs - Another warp (or the same warp) handles all vector outputs: 32 threads, each thread does 1 output To make memory coalesced: - Instead of thread `t` processing output `t`, I should have all threads in the warp work on the SAME output at a time, with each thread handling a subset of the inputs. So for output `o` and item `item`: - Warp loads all 80 combined values into shared memory or registers - Each thread of the warp computes `combined[i] * weight[o][i]` for a subset of `i` - Then do warp reduction to sum This way, at each step all threads access the same weight row but different columns. Still not coalesced because columns are accessed by different threads. Wait, that's coalesced! If thread 0 accesses weight[0], thread 1 accesses weight[1], etc., that's adjacent in memory and is coalesced! But the combined values are also accessed with each thread reading different `i`, which is also coalesced. And after computing partial sums, we do a warp shuffle to accumulate. Let me implement this: ```cpp // Each warp processes one output for one item // Within a warp: thread t processes elements t, t+32, t+64, ... of the input // Then warp reduction float sum = 0.0f; for (int i = tid; i < SI + VI; i += 32) { sum += combined[i] * weight_row[i]; } // Warp reduction for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xffffffff, sum, offset); } if (tid == 0) { s_out[o] = sum + bias; } ``` Now, if I have 32 threads per warp and one warp per output: - For SO=64: need 64 warps for all outputs of one item - That's 2048 threads per item. With 1024 items, that's just too many. I need to process multiple items or have fewer warps per item. What if each warp processes all outputs (SO + VO) for one item? That's: - SO=64, VI0=16 - A warp has 32 threads. - For scalar: each thread processes all 80 inputs but only for 2 outputs. Not great for memory reuse. Actually, if combined and vectors are already loaded into shared memory (which is fast), then the weight matrix is the main bottleneck. The weight matrices are small: - scalar_weight: 64 * 80 = 5120 floats = 20KB - This can fit in L1 cache or be read multiple times efficiently. Let me try a different approach: - One thread per output (for scalar and vector combined) - Process the same output across multiple items in a thread Actually, the best approach for small matrices is often to transpose the weights and load them into registers, or use shared memory for weights. Given the sizes are very small (80 inputs, 64 + 16 outputs), let me try: 1. Load all weights into shared memory 2. Each item is processed by one thread that loops over all outputs Items = 1024, threads = 1024, blocks = 4 (256 threads each). This has very few blocks. Items = 1024, threads = 128, blocks = 8. Still few blocks. Actually wait, the kernel launch overhead is small. What matters is thread-level parallelism. With 8 blocks * 128 threads = 1024 threads, we'll have low occupancy per SM. Each SM can run many more threads. What about having one block handle one output for all items? Then we'd have: - SO + VO = 80 outputs - 80 blocks - Each block has 256 threads handling 1024 items -> 4 items per thread Or better: - One block handles all outputs for a group of items - Group size = block size = 256 threads = 256 items - Number of blocks = 1024 / 256 = 4 Still only 4 blocks. OK, let me just try to maximize warp-level parallelism. With 1024 total items to process, and we want say 128 blocks, each block handles 8 items. So if each block has 256 threads, there are 32 threads per item. For each item, 32 threads cooperate: - First, load combined features and vectors into shared memory - Then, for scalar outputs: 32 threads handle all 64 outputs, 2 per thread - For vector outputs: 32 threads handle all 16 outputs, 1 per thread (with all 3 components) To improve coalescing for weight access: - Have all threads loop over the same output `o`, with each thread computing partial sums for different `i` values, then warp reduce. - With this, each warp processes one output. - 32 threads can handle 32 different outputs across the 32 warps. Wait, a block of 256 threads = 8 warps. If each warp handles one output: - 8 warps handle 8 outputs at a time - For 64 scalar outputs: 64/8 = 8 iterations - For 16 vector outputs: 16/8 = 2 iterations Each warp processes its assigned output: ```cpp int warp_id = threadIdx.x / 32; int lane_id = threadIdx.x % 32; for (int o_iter = 0; o_iter < num_scalar_iters; o_iter++) { int o = warp_id + o_iter * 8; if (o >= SO) continue; float partial = 0; for (int i = lane_id; i < SI + VI; i += 32) { partial += combined[i] * weight[o][i]; } // warp reduce for (int offset = 16; offset > 0; offset >>= 1) { partial += __shfl_down_sync(0xffffffff, partial, offset); } if (lane_id == 0) s_out[o] = partial + bias[o]; } ``` This is much better for memory coalescing because: - At step i, all threads in the warp access weight[o][i], weight[o][i+32], etc. -> NOT coalesced, they all access the same address within the warp. Wait, that's actually a broadcast, not coalesced. But broadcasts are also efficient. For the combined array, each thread accesses combined[i], combined[i+32], etc. -> NOT coalesced across the warp, but they access the same shared memory bank? Actually, shared memory access doesn't need coalescing, bank conflicts are the issue. If threads access consecutive elements (i, i+32, i+64), that's 32 banks apart (mod 32), so no bank conflicts! For weights, since all threads access `weight[o][i]`, `weight[o][i+32]`, etc., this is sequential access and coalesced. OK so this approach should be good. Let me implement it. Wait, I also need to handle multiple items per block. Let me say: - Items per block = 8 - Each item gets 32 threads (1 warp) - Total threads per block = 8 * 32 = 256 - Number of blocks = 1024 / 8 = 128 Actually, I can use fewer items per block to reduce shared memory pressure, but for these small arrays it's fine. Let me rewrite the kernel with this warp-based approach: ```cpp __global__ void fused_se3_kernel(...) { const int items_per_block = 8; const int warps_per_item = 1; const int threads_per_item = 32; int item_local = threadIdx.x / threads_per_item; int lane = threadIdx.x % threads_per_item; int item_idx = blockIdx.x * items_per_block + item_local; if (item_idx >= B * N) return; extern __shared__ float smem[]; float* combined = smem + item_local * (SI + VI); float* vectors = smem + items_per_block * (SI + VI) + item_local * (VI * 3); // Load data... float* s_out_ptr = scalar_out + item_idx * SO; float* v_out_ptr = vector_out + item_idx * (VO * 3); // Scalar outputs: each warp handles one output // Wait, but there are only 32 threads per item and 64 outputs. // We need multiple iterations. // Actually, let me not tie warp to output. Instead, for each output, // all threads in the warp compute partial sums and reduce. for (int o = lane; o < SO; o += 32) { float partial = 0.0f; #pragma unroll for (int i = 0; i < SI + VI; ++i) { partial += combined[i] * scalar_weight[o * (SI + VI) + i]; } // No need for reduction if each thread handles one output s_out_ptr[o] = partial + scalar_bias[o]; } // Same for vector... } ``` Wait, this is what I had before, but the weight access pattern is bad. Let me try the partial sum + warp reduction approach: ```cpp __global__ void fused_se3_kernel(...) { int item_idx = blockIdx.x; int lane = threadIdx.x; int warp_id = threadIdx.x / 32; // 0..7 for 256 threads // 256 threads handle 8 items? No, let's say 1 item per block with 256 threads // blocks = 1024, threads = 256 // But we can process 8 outputs simultaneously with 8 warps: extern __shared__ float smem[]; float* combined = smem; float* vectors = smem + (SI + VI); // Load... int warp_id = threadIdx.x / 32; int lane = threadIdx.x % 32; // 8 warps, each handles one output at a time int num_scalar_iters = (SO + 7) / 8; for (int iter = 0; iter < num_scalar_iters; iter++) { int o = iter * 8 + warp_id; if (o >= SO) continue; float partial = 0.0f; #pragma unroll for (int i = lane; i < SI + VI; i += 32) { partial += combined[i] * scalar_weight[o * (SI + VI) + i]; } // Warp reduction #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { partial += __shfl_down_sync(0xffffffff, partial, offset); } if (lane == 0) { s_out_ptr[o] = partial + scalar_bias[o]; } } // Similar for vector... } ``` With 1024 blocks and 256 threads each, we get 262144 threads total. That's excellent parallelism. But wait, each block only processes 1 item, so we need 1024 blocks. With 256 threads per block, that's 262K threads. For each block: - 8 warps process 8 outputs simultaneously - For SO=64, 8 iterations - For VO=16, 2 iterations Each warp's memory access to weights: at step i, warp accesses `scalar_weight[o][i]`, `scalar_weight[o][i+32]`, etc. Since all threads in the warp access the same `o` but different `i`, this is strided access. With a row of 80 floats, and threads accessing `i, i+32, i+64`, thread 0 accesses index 0, thread 1 accesses index 1, etc. up to thread 31 accessing index 31. Then thread 0 accesses index 32, etc. Actually, since lane_id maps to consecutive `i` values (0-31, 32-63, 64-79), this IS coalesced! Threads 0-31 access consecutive memory elements. The only issue: the row is 80 floats. When accessing elements 0-31, we get coalesced access. 32-63: also coalesced. 64-79: only 16 threads do work here, accessing elements 64-79, also coalesced. So this is great! The combined array is in shared memory, so no coalescing concerns there. And we process 8 outputs (8 warps) in parallel, giving good utilization. For vector outputs: - 2 iterations (VO=16, 8 warps per iteration) - Each output has 3 components, computed as 3 partial sums - After partial sums are reduced, multiply by gate value But wait, computing the gate value also requires the same combined array. We can compute all gate values first, store them in shared memory, then use them for the vector computation. Or, since the combined array is already in shared memory, each warp can compute its gate value as part of the vector output loop. Let me write this kernel. Also, since we have 1024 blocks, each with shared memory for combined (80 floats) + vectors (48 floats) = 512 bytes + maybe gate values (16 floats) = 576 bytes, this is very small and shouldn't limit occupancy. Let me also add `#pragma unroll` with actual unroll counts. Since SI+VI=80, I can unroll to some extent. Actually, for 80 iterations, `#pragma unroll` might increase register pressure. Let me be selective: - Unroll the inner vector component loop (3 iterations) - Unroll the warp reduction (fixed 5 iterations for 32 threads) Let me also try using `__ldg` (load global) for weight access to leverage read-only cache, or mark weights as `__restrict__ const` and let the compiler handle it. Here's my new kernel plan: - 1024 blocks, each handling 1 item - 256 threads per block (8 warps) - Combined and vectors loaded into shared memory by all threads - For scalar outputs: 8 warps cooperate, each warp handles one output, iterating 8 times for SO=64 - For vector outputs: similar, 2 iterations for VO=16 - Each warp computes partial sums with warp reduction Actually, let me reconsider. 1024 blocks * 256 threads = 262k threads. But we're only launching one wave actually. H100 has 128 SMs and each SM can run multiple blocks concurrently. Actually, the main concern is just getting enough parallelism. 1024 blocks is more than enough for good scheduling. Let me write this kernel now. I want to make sure the math is correct. ```cpp __global__ void fused_se3_kernel( const float* __restrict__ scalar_feat, const float* __restrict__ vector_feat, const float* __restrict__ scalar_weight, const float* __restrict__ scalar_bias, const float* __restrict__ vector_weight, const float* __restrict__ gate_weight, const float* __restrict__ gate_bias, float* __restrict__ scalar_out, float* __restrict__ vector_out, int B, int N, int SI, int SO, int VI, int VO ) { const int item_idx = blockIdx.x; if (item_idx >= B * N) return; const int lane = threadIdx.x % 32; const int warp_id = threadIdx.x / 32; const int warps_per_block = blockDim.x / 32; extern __shared__ float smem[]; float* combined = smem; float* vectors = smem + (SI + VI); float* gate_vals = vectors + (VI * 3); // Load scalar features const float* s_in = scalar_feat + item_idx * SI; for (int i = threadIdx.x; i < SI; i += blockDim.x) { combined[i] = s_in[i]; } // Load vector features const float* v_in = vector_feat + item_idx * (VI * 3); for (int i = threadIdx.x; i < VI * 3; i += blockDim.x) { vectors[i] = v_in[i]; } __syncthreads(); // Compute norms for (int i = threadIdx.x; i < VI; i += blockDim.x) { float v0 = vectors[i * 3 + 0]; float v1 = vectors[i * 3 + 1]; float v2 = vectors[i * 3 + 2]; combined[SI + i] = sqrtf(v0 * v0 + v1 * v1 + v2 * v2); } __syncthreads(); float* s_out_ptr = scalar_out + item_idx * SO; float* v_out_ptr = vector_out + item_idx * (VO * 3); // Scalar linear outputs const int scalar_iters = (SO + warps_per_block - 1) / warps_per_block; for (int iter = 0; iter < scalar_iters; iter++) { int o = iter * warps_per_block + warp_id; if (o >= SO) continue; float partial = 0.0f; const float* w_row = scalar_weight + o * (SI + VI); for (int i = lane; i < SI + VI; i += 32) { partial += combined[i] * w_row[i]; } #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { partial += __shfl_down_sync(0xffffffff, partial, offset); } if (lane == 0) { s_out_ptr[o] = partial + scalar_bias[o]; } } // Precompute gate values const int gate_iters = (VO + warps_per_block - 1) / warps_per_block; for (int iter = 0; iter < gate_iters; iter++) { int o = iter * warps_per_block + warp_id; if (o >= VO) continue; float partial = 0.0f; const float* gw_row = gate_weight + o * (SI + VI); for (int i = lane; i < SI + VI; i += 32) { partial += combined[i] * gw_row[i]; } #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { partial += __shfl_down_sync(0xffffffff, partial, offset); } if (lane == 0) { gate_vals[o] = 1.0f / (1.0f + expf(-(partial + gate_bias[o]))); } } __syncthreads(); // Vector outputs for (int iter = 0; iter < gate_iters; iter++) { int o = iter * warps_per_block + warp_id; if (o >= VO) continue; float g = gate_vals[o]; const float* vw_row = vector_weight + o * VI; float vsum0 = 0.0f, vsum1 = 0.0f, vsum2 = 0.0f; for (int i = lane; i < VI; i += 32) { float w = vw_row[i]; vsum0 += vectors[i * 3 + 0] * w; vsum1 += vectors[i * 3 + 1] * w; vsum2 += vectors[i * 3 + 2] * w; } vsum0 = __shfl_down_sync(0xffffffff, vsum0, 16) + vsum0; vsum0 = __shfl_down_sync(0xffffffff, vsum0, 8) + vsum0; vsum0 = __shfl_down_sync(0xffffffff, vsum0, 4) + vsum0; vsum0 = __shfl_down_sync(0xffffffff, vsum0, 2) + vsum0; vsum0 = __shfl_down_sync(0xffffffff, vsum0, 1) + vsum0; vsum1 = __shfl_down_sync(0xffffffff, vsum1, 16) + vsum1; vsum1 = __shfl_down_sync(0xffffffff, vsum1, 8) + vsum1; vsum1 = __shfl_down_sync(0xffffffff, vsum1, 4) + vsum1; vsum1 = __shfl_down_sync(0xffffffff, vsum1, 2) + vsum1; vsum1 = __shfl_down_sync(0xffffffff, vsum1, 1) + vsum1; vsum2 = __shfl_down_sync(0xffffffff, vsum2, 16) + vsum2; vsum2 = __shfl_down_sync(0xffffffff, vsum2, 8) + vsum2; vsum2 = __shfl_down_sync(0xffffffff, vsum2, 4) + vsum2; vsum2 = __shfl_down_sync(0xffffffff, vsum2, 2) + vsum2; vsum2 = __shfl_down_sync(0xffffffff, vsum2, 1) + vsum2; if (lane == 0) { v_out_ptr[o * 3 + 0] = g * vsum0; v_out_ptr[o * 3 + 1] = g * vsum1; v_out_ptr[o * 3 + 2] = g * vsum2; } } } ``` For the warp reduction of vsum, since VI=16 and each warp has 32 threads, only threads 0-15 do work. We need to reduce across the 16 active threads. We can use a mask approach or just do the full 32-thread reduction. Actually, the simpler approach: ```cpp // First, put the sum from thread i into thread i+16 float vsum0_16 = __shfl_down_sync(0xffffffff, vsum0, 16); vsum0 += (lane < 16) ? vsum0_16 : 0; // Then reduce the lower 16 threads for (int offset = 8; offset > 0; offset >>= 1) { vsum0 += __shfl_down_sync(0xffff, vsum0, offset); } // Now lane 0 has the full sum ``` Or even simpler, since threads 16-31 have vsum0=0, the full 32-thread reduction works correctly. Let me also think about shared memory size. For one item: - combined: 80 floats - vectors: 48 floats - gate_vals: 16 floats Total: 144 floats = 576 bytes With 256 threads per block, that's negligible. Let me write this kernel and test it. For the reduction of vsum, let me use a helper function if C++ allows it, or just inline it. Since load_inline compiles as C++, I can write a device function.

The kernel is correct. Let me also run a static check to make sure there are no issues before submitting. The current kernel uses: - One block per (B, N) item - Fuses all operations: vec norm, scalar linear, gate sigmoid, vector linear in a single kernel - Uses shared memory for combined features and vectors - Uses warp-level reductions This should be efficient. Let me do a quick static check first, then submit.

The static check passed and the kernel is correct. Let me submit my best kernel - the fused kernel with warp-level reductions (v2) which had good correctness and should be faster than the reference due to: 1. Fusing all operations into a single kernel 2. Reducing memory traffic by using shared memory 3. Using warp-level reductions for efficient parallel reduction 4. Avoiding intermediate cat/norm/einsum/sigmoid operations that create extra memory allocations in PyTorch