We need to design a custom CUDA kernel that computes the Outer-Product Mean operation described. The reference does: 1. LayerNorm over msa_dim. 2. Two linear projections proj_a and proj_b (msa_dim -> hidden_dim). 3. Masking a and b by (B,S,N,1) mask. 4. Outer product: einsum("bsih,bsjk->bijhk", a, b) resulting (B,N,N,H,H). 5. denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1) 6. Divide outer by denom. 7. Reshape outer to (B,N,N,H*H) 8. Linear out_proj (H*H -> pair_dim). Goal: fuse as much as possible into a single or few CUDA kernels, improving speed over PyTorch's many kernel launches and intermediate allocations. We need to use torch.utils.cpp_extension.load_inline. We need to produce numerically equivalent output at fp32 precision. (The tolerance for fp32 is atol=rtol=1e-4.) We can also consider doing the layer norm and linear projections in PyTorch (they map to efficient cuBLAS / LayerNorm kernels) and fuse only the outer product mean, masking, division, reshape, and final linear projection. However, we might also try to fuse the entire computation into one big kernel to reduce memory traffic. Let's analyze arithmetic intensity. Dimensions: batch_size = 2 msa_depth = 8 seq_len = 32 msa_dim = 64 hidden_dim = 32 pair_dim = 64 So N=32, S=8, H=32. LayerNorm: B*S*N = 2*8*32 = 512 rows of length 64. Linear proj_a: (512,64) @ (64,32) => small gemm. Linear proj_b: same. Masking: elementwise. Outer product: for each batch b, for each i in N, j in N, compute sum over s of a[b,s,i,h1] * b[b,s,j,h2]. Since S=8, it's a small reduction. Denominator: sum over s of msa_mask[b,s,i] (actually mask sum per position i? Wait denom is sum over dim=1 of mask, shape (B,N). It doesn't depend on j. They unsqueeze to (B,N,1,1,1) and broadcast. So denominator depends only on position i (or maybe on N dimension) but not on j. Wait they sum over dim=1 (msa_depth), producing (B,N). So denominator for pair (i,j) is sum_s mask[b,s,i]. Actually they don't sum over positions; they sum over MSA rows, but the mask shape is (B,S,N). Summing over dim=1 gives (B,N). They then unsqueeze to (B,N,1,1,1). That seems asymmetric: denominator depends only on i, not j. But maybe the mask is the same across N? In typical AF2, msa_mask shape (B,S) I think. But here they have shape (B,S,N) and sum over S. So denom[b,i] = sum_s mask[b,s,i]. If mask is all ones, denom = S = 8. Thus division is a broadcast over j and h1,h2. Then reshape and out_proj: (B,N,N,H*H) @ (H*H, pair_dim). This is a GEMM on each (B,N,N) slice? Actually reshape to (B*N*N, H*H) and matmul with (H*H, pair_dim). Since B*N*N = 2*32*32 = 2048, H*H=1024, pair_dim=64. So it's a GEMM of size 2048x1024 @ 1024x64. We could fuse the entire thing into a single kernel that, for each output pair element (b,i,j,d) computes the weighted sum over h1,h2 of out_proj_weight[d, h1*H+h2] * (1/denom) * sum_s a[b,s,i,h1] * b[b,s,j,h2]. This would avoid materializing the (B,N,N,H,H) tensor (~2*32*32*32*32 = ~2M floats, 8MB). But we can also compute intermediate (B,N,N,H,H) using shared memory. Since N=32, H=32, B=2. But the output is 2*32*32*64 = 131072 floats. However, evaluating all possible fusions might be tricky but doable. Simpler approach: compute a and b with PyTorch (layer norm + linear), then write a fused kernel for outer product mean + out_proj. Or compute outer product mean with a custom kernel that directly outputs (B,N,N,pair_dim) by contracting over s, h1, h2. Let's think about performance. The reference likely does: - LayerNorm kernel. - Two linear layers (cuBLAS). - Elementwise multiply for mask. - einsum (likely uses bmm or tensordot; but BS,N,H times BS,N,H -> B,N,N,H,H or similar). In PyTorch, if einsum is not optimized, it may expand and do a lot of ops. - division - reshape, linear. The einsum is expensive because it creates a large intermediate and does many operations: B*N*N*H*H*S = 2*32*32*32*32*8 = 16,777,216 multiply-adds. That's not huge. But memory bandwidth for writing 8MB intermediate and reading it again might matter. If we fuse the einsum+division+reshape+out_proj into a single kernel, we can reduce memory bandwidth by reading a and b once and writing final output directly. The arithmetic is about ~16M FLOPs for the outer product + division, plus ~2M * 1024 * 64? Wait the final linear maps each (H*H) vector to pair_dim. That's B*N*N*H*H*pair_dim = 2*32*32*1024*64 = 134,217,728 FLOPs. That's actually 134M FLOPs. Combined with 16M for outer product. Total ~150M FLOPs. That's not huge but memory bound if doing many separate kernels. If we fuse the final linear into the outer product kernel, we compute for each (b,i,j,d) the dot product of out_proj_weight[d,:] with the normalized outer product vector. This is basically a small GEMV per (b,i,j). We could tile over H and pair_dim. Alternatively, we can compute the outer product mean intermediate ((B,N,N,H*H)) with a fused kernel and then call a standard linear (GEMM). That intermediate is 8MB. Writing 8MB and reading it plus reading weight and writing output to GEMM. The GEMM is 2048x1024 @ 1024x64, which is tiny and might not saturate GPU. So fusing the final linear is beneficial. But we also have LayerNorm and the two linears. We can keep those in PyTorch. The total memory traffic for a,b is small: a and b each are B*S*N*H = 2*8*32*32 = 16384 floats (~64KB). So reading them is negligible. Thus, the major bottleneck is the einsum + division + out_proj. Fusing them into a single kernel is the best. To write this kernel, we will have a CUDA kernel launched with blocks over (B, N, N/ something) and threads over pair_dim and maybe H tiles. Given N=32, B=2, we have 2048 (b,i,j) positions. For each, we compute a vector of length pair_dim (64). Each output element d = dot over h1,h2 of W[d, h1*H+h2] * (sum_s A[s,i,h1] * B[s,j,h2] / denom_i). We can precompute for each pair (i,j) the outer product vector O[h1*H+h2] = sum_s A[s,i,h1] * B[s,j,h2] / denom_i. Then out = W @ O. This is GEMV of size 64 x 1024 per position. Since there are 2048 positions, total work is 2048 * 64 * 1024 = 134M mul-adds. At 150M FLOPs, the kernel is compute-bound if we keep data in registers/shared memory. But maybe memory bound if we load W repeatedly. However W is small: 64*1024 = 65536 floats (~256KB). It can be cached in L2. A kernel can have each block compute all pair_dim outputs for a tile of (i,j) positions. For example, each block handles one (b, i, j) and uses 64 threads to compute the 64 output elements using parallel reduction over H*H. But 2048 blocks of 64 threads = 131072 threads. That is a decent amount, but each block has very little parallelism (64 threads) and the reduction over 1024 elements per thread might be okay: each thread computes one output d, loop over h1,h2 reading W[d, idx] and O[idx]. But O[idx] depends on s,h1,h2. We can compute O[idx] on the fly: for each output d, thread sums over h1,h2 W[d, h1*H+h2] * (sum_s A[s,i,h1] * B[s,j,h2] / denom_i). If we reorder loops, we can compute sum_s A[s,i,h1] * B[s,j,h2] in shared memory or registers first, but maybe better to directly accumulate into output. Simpler: For each (b,i,j), we can load a_chunk = A[b, :, i, :] (S=8, H=32) and b_chunk = B[b, :, j, :] into shared memory. Then for each pair (h1,h2), compute o = sum_s a[s,h1] * b[s,h2] / denom. Then each thread (for each d) accumulates W[d, h1*H+h2] * o. This requires shared memory for a_chunk (8*32=256 floats) and b_chunk (256 floats) and maybe for W tile? W is 64*1024. If we have 64 threads, each thread can load one column of W? Actually we need W per d. Each thread owns one d. So each thread would read 1024 floats from W. But W is in global memory; L2 cache may help. 2048 blocks * 64 threads each reading 1024 floats from W = 134M loads. W is small enough to fit in L2 shared among SMs. Alternatively, tile W in shared memory: load a tile of W of size 64 x tile_size into shared memory? Since pair_dim=64, all threads (64) could cooperatively load a tile of W of size 64 x 32 for example. Then iterate over h2 tile. But maybe a simpler approach: write a kernel that computes the outer product intermediate (B,N,N,H,H) in one kernel, store it, then call linear. But storing 8MB intermediate might be okay but could be memory bound. However we have to beat PyTorch reference. The reference is likely not heavily optimized. Let me think: PyTorch einsum might be doing a lot of memory traffic. If we replace just einsum + division + reshape + out_proj with a single custom kernel that computes output directly, we can likely see a good speedup. Given the small sizes, we can also write a kernel that uses shared memory explicitly to compute the outer product and then multiplies by weight matrix on the fly. Let's consider a block configuration. Option A: one output element per thread. Each thread handles one (b,i,j,d) and loops over s, h1, h2: acc = 0 for h1 in 0..H-1: for h2 in 0..H-1: o = 0 for s in 0..S-1: o += a[b,s,i,h1] * b[b,s,j,h2] o /= denom[b,i] acc += W[d, h1*H+h2] * o This is 1024 * 8 multiply-adds per thread = 8192 iterations. Each iteration reads a and b and W. But W is read per d per idx. That's a lot of reads (134M*1024? Wait no, each thread reads W[d, idx] once per idx. Total W reads = num_outputs * H*H = 2048*64*1024 = 134M reads. That's okay. Total a reads = for each output, each thread reads a[b,s,i,h1] for each h1,h2,s. That's 2048*64 * 8 * 32 * 32 = 33,554,432 reads per variable? Wait 2048*64 = total output elements (131072). Each output element loops over 1024 * 8 = 8192 multiplies. But a and b are reused across d. For fixed (b,i,j), all d threads read same a and b values. So within a block, we can share a and b via shared memory. But across different blocks for same (b,i,j) but different d, we duplicate reads. If we schedule one (b,i,j) per block and 64 threads for d, then a and b are loaded once per block. Option B: one block per (b,i,j). Threads = pair_dim (64). Shared memory: load a_tile (S, H) and b_tile (S, H). Since S=8, H=32, total floats = 256+256=512, 2KB. Good. Within block: - Load a_tile and b_tile into shared memory. - Compute denominator d_i. - For each h1 in 0..H-1: For each h2 in 0..H-1: Compute o = sum_{s=0}^{S-1} a_tile[s][h1] * b_tile[s][h2] / denom. Each thread (d) reads W[d][h1*H+h2] and accumulates W * o. We can have threads cooperatively compute o for each (h1,h2) and store in a register? But 64 threads can't compute 1024 o's simultaneously. We can loop over h1 and h2. For each (h1,h2), one thread could compute the sum over s? That's 8 multiplies, trivial. Then we broadcast to all threads. But 1024 iterations per block, each with a small sum. Not too bad. Alternatively, we can restructure: For each h1, precompute for all h2 the sum over s in registers. Since S=8, H=32, we can have each thread compute a partial product. Actually we can use shared memory for a and b. Implementation details: We can define BLOCK_D = 64 (pair_dim). launch kernel with grid_dim = (B, N, N). block_dim = (BLOCK_D,). So each block handles one (b,i,j). 2048 blocks total. Shared memory: - float a_shared[S * H]; - float b_shared[S * H]; Loading: Each thread loads some elements. Total elements = S*H = 256. With 64 threads, each loads 4 elements. Compute denom = sum_s a_mask? Actually we need denom_i = sum_s mask[b,s,i]. Wait we only have a and b after masking. In the reference, a and b are multiplied by mask, and then denominator is sum of mask over S. So a and b are already masked. But wait: if we compute a = proj_a(x) * mask, then a already contains zeros for masked entries. The denominator is clamp(sum_s mask_i, 1). Since a and b are zeroed, the sum over s of a[s,i,h1] * b[s,j,h2] effectively only includes valid s. But dividing by denom_i (which is number of valid s for position i) gives mean over valid s. However note that denominator uses only mask, not the a and b values. So if we have a and b already zeroed, we need to also pass the denominator values or the mask to compute denom. We can compute denom on CPU easily: denom = msa_mask.sum(dim=1) which yields shape (B,N). So we can compute denom in PyTorch and pass it to the kernel along with a and b. Or pass the mask to the kernel. Since B, N, S are small, computing denom on CPU and passing as a tensor is fine. Wait, but if we fuse the linear projections into the kernel too, we could compute denom earlier. But easier: in ModelNew.forward, compute x = layer_norm(msa_repr), a = proj_a(x), b = proj_b(x), denom = msa_mask.sum(dim=1). Then call the kernel with a, b, denom, out_proj weight and bias. But we also need to mask a and b. In the reference, a = a * mask, b = b * mask. Since mask is (B,S,N,1), we can simply do a = a * mask; b = b * mask in PyTorch before calling the kernel. The small elementwise multiply is cheap. Or we could incorporate masking into the kernel by passing mask. But elementwise multiply is trivial. Thus, the custom kernel inputs: a (B,S,N,H), b (B,S,N,H), denom (B,N), weight (pair_dim, H*H), bias (pair_dim). The kernel computes for each (b,i,j,d): out = bias[d] + sum_{h1,h2} weight[d, h1*H+h2] * (sum_s a[b,s,i,h1] * b[b,s,j,h2] / denom[b,i]) But wait: denominator uses clamp(min=1). So we need to ensure denom >= 1. We can do denom = denom.clamp(min=1) on CPU. Also note: denom_i depends only on i, not j. So the kernel can read denom_i once per block. Potential issue: if denom_i is zero (all masked), reference clamps to 1, so division yields 0/1 = 0. Since a and b are zeroed, sum_s a*b = 0. So result is zero. Using denom=1 yields zero. Good. Now, about numerical precision: we need to compute in float (fp32). Using float throughout should match. Let's design the kernel. Pseudo CUDA: __global__ void outer_product_mean_proj_kernel( const float* __restrict__ a, const float* __restrict__ b, const float* __restrict__ denom, const float* __restrict__ weight, const float* __restrict__ bias, float* __restrict__ out, int B, int S, int N, int H, int pair_dim) { // blockIdx.x = b // blockIdx.y = i // blockIdx.z = j // threadIdx.x = d (0..pair_dim-1) int b = blockIdx.x; int i = blockIdx.y; int j = blockIdx.z; int d = threadIdx.x; // Load a and b for this b,i and b,j into shared memory __shared__ float a_shared[8*32]; // S*H __shared__ float b_shared[8*32]; int base_a = ((b * S * N) + i) * H; // a[b, :, i, :] contiguous? Wait shape is (B,S,N,H). Strides: a.stride(0)=S*N*H, a.stride(1)=N*H, a.stride(2)=H, a.stride(3)=1. // But we don't know if it's contiguous. Better pass strides. } Hmm, we need to be careful with tensor strides for non-contiguous memory, but PyTorch outputs from linear are contiguous (B,S,N,H). However denom is (B,N) contiguous. Weight is (pair_dim, H*H) contiguous if we transpose? Actually in PyTorch linear weight is (out_features, in_features). So out_proj.weight shape is (pair_dim, H*H), row-major (outer dimension pair_dim). So weight[d, :] is contiguous. To keep it simple, we can enforce contiguous with .contiguous() in Python before passing pointers. Thus assume a, b are contiguous with shape (B, S, N, H). Strides: a[b,s,i,h] offset = ((b * S + s) * N + i) * H + h. Denom shape (B, N), offset = b*N + i. So: int bs_idx = blockIdx.x; int i = blockIdx.y; int j = blockIdx.z; int d = threadIdx.x; int a_base = ((bs_idx * S) * N + i) * H; // a[bs_idx, 0, i, 0] int b_base = ((bs_idx * S) * N + j) * H; // a and b have same shape. float denom_val = denom[bs_idx * N + i]; // denom_val >= 1 // cooperatively load a and b into shared memory int total = S * H; for (int idx = d; idx < total; idx += blockDim.x) { int s = idx / H; int h = idx % H; a_shared[idx] = a[a_base + s * N * H + h]; b_shared[idx] = b[b_base + s * N * H + h]; } __syncthreads(); float acc = 0.0f; for (int h1 = 0; h1 < H; ++h1) { for (int h2 = 0; h2 < H; ++h2) { float sum_s = 0.0f; for (int s = 0; s < S; ++s) { sum_s += a_shared[s * H + h1] * b_shared[s * H + h2]; } sum_s /= denom_val; int w_idx = d * (H * H) + h1 * H + h2; acc += weight[w_idx] * sum_s; } } if (bias != nullptr) acc += bias[d]; int out_base = ((bs_idx * N + i) * N + j) * pair_dim + d; out[out_base] = acc; This kernel has each thread loop over H*H*S = 8*1024 = 8192 operations. 64 threads per block => 524k ops per block. With 2048 blocks, total ~1.07B ops? Wait 2048 blocks * 64 threads * 8192 = 1,073,741,824 operations (multiplies + adds). That's about 1B FLOPs. At 1 TFLOPS, that's 1 ms. Memory bandwidth for reading a,b per block: 2 * 256 * 4 = 2KB. 2048 blocks => 4MB read. Weight reads: each block reads 1024 floats per thread * 64 threads * 4 bytes = 256KB per block. 2048 blocks => 512MB read. That's more traffic. But weight is cached. L2 cache on GPU is maybe 4-6MB. Weight is 256KB. So reading weight 2048 times might be okay if it's L2 cached. But each block reads the entire weight matrix. However weight is small and constant across (i,j) but varies across b? No, weight is same for all batches and positions. So it will be heavily cached. The L2 can hold it, so reads are L2 hits. That should be fast. But each thread reads weight[d, :] which is contiguous (stride 1) of length 1024. In L2, weight is stored. Coalesced across threads? For a fixed (h1*H+h2), threads read weights for different d, so memory access is strided by H*H = 1024. That's not coalesced! Wait: weight is stored as [pair_dim, H*H]. So weight[d, idx] address = d * (H*H) + idx. For a fixed idx, threads with consecutive d read memory spaced by 1024 floats. This is terrible for coalescing. They will be loaded from separate cache lines. We can transpose the weight before sending to kernel, so that we can read weight[idx, d] i.e. for each idx, the d dimension is contiguous. Or we can load weight into shared memory per block? But weight is 64x1024 floats = 256KB. Shared memory per SM is limited (164KB on A100). So we cannot load entire weight into shared memory if we also need a_shared and b_shared. But we don't need to load all of it; we can tile. Alternatively, we can restructure the accumulation to read weight in a coalesced way. For each (h1,h2) we need weight[:, h1*H+h2] for all d. If weight is stored transposed (H*H, pair_dim), then this column is contiguous. So we should pass weight as (H*H, pair_dim) contiguous. We can do weight_t = out_proj.weight.t().contiguous() in Python. Then w_idx = (h1*H+h2) * pair_dim + d. Threads with consecutive d read consecutive addresses. That is fully coalesced. Great. So in Python: W = self.out_proj.weight.t().contiguous() # shape (H*H, pair_dim) bias = self.out_proj.bias # shape (pair_dim,) Then in kernel: float w_val = weight[(h1*H + h2) * pair_dim + d]; Because pair_dim = blockDim.x = 64. Then accumulate: acc += w_val * sum_s; This is coalesced. Now, about loops: For each (h1,h2), one thread computes sum_s. But that sum_s is same for all threads (same h1,h2). So all 64 threads compute the exact same sum_s in their local registers. That's redundant. We could compute sum_s once and broadcast to all threads via shared memory. But doing it redundantly means each thread does S=8 multiplies per (h1,h2). With H*H=1024, that's 8192 multiplies per thread. 64 threads * 8192 = 524k operations per block. That's okay. Redundant computation but saves shared memory. However, we can have threads compute sum_s in parallel and store in shared memory. Better: use a small shared memory for the outer product chunk. For each (h1,h2), we need O = sum_s a_s_h1 * b_s_h2. That's a scalar. If we have 64 threads, we can assign each thread to compute one O element for a subset of (h1,h2). But we have 1024 O elements. 64 threads can compute 16 each in parallel, but then we need to store them. 1024 floats = 4KB, which fits in shared memory. Then we can do a matrix-vector multiply: weight_t (1024 x 64) @ O (1024) -> out (64). We can have threads compute dot product of weight_t column d with O. That's exactly one output per thread. We can load O into shared memory, then each thread loops over 1024 elements reading weight_t and O. This eliminates redundant inner product over s. So steps per block: 1. Load a_tile (S*H=256 floats) and b_tile (256 floats) into shared memory. 2. Compute O[h1*H+h2] = sum_s a_shared[s,h1] * b_shared[s,h2] / denom. We have 64 threads. For idx = threadIdx.x; idx < H*H; idx += blockDim.x: h1 = idx / H; h2 = idx % H; float sum = 0; for s in 0..S-1: sum += a_shared[s*H+h1] * b_shared[s*H+h2]; O_shared[idx] = sum / denom_val; __syncthreads(); 3. Now each thread (d) computes acc = bias[d]; for idx in 0..H*H-1: acc += weight[idx * pair_dim + d] * O_shared[idx]; out = acc. This requires O_shared of size H*H = 1024 floats = 4KB. a_shared + b_shared = 2KB. Total shared memory = 6KB per block. That's fine. Total FLOPs per block: Step 2: 1024 * 8 = 8192 mul-adds, done by 64 threads => each thread does 128 * 8 = 1024 ops, plus some loops. Step 3: 1024 * 64 = 65536 mul-adds. Each thread does 1024 ops. Total ~72k ops per block. 2048 blocks => ~147M ops. That matches earlier estimate. Memory reads: - a_tile: 1KB - b_tile: 1KB - weight: 1024*64*4 = 256KB per block? Wait each block reads the entire weight matrix (256KB). 2048 blocks => 512MB. But weight is small and cached. Reads of weight are coalesced: for each idx, threads read consecutive floats weight[idx*pair_dim + d]. This is coalesced within a warp. Since we loop idx over 0..1023, each warp reads pair_dim floats per idx. That's 64 floats per warp = 256 bytes per idx. Over 32 warps? Actually block has 64 threads = 2 warps. For each idx, warp 0 reads idx*64 ... idx*64+31, warp 1 reads idx*64+32 ... idx*64+63. Those are contiguous 256 bytes. Coalesced into 4 64-byte transactions (or 2 128-byte). Good. - O_shared is in shared memory, no global load aside from storing. So this kernel should be quite fast. We need to pass bias as well. Implementation in CUDA: template <int S, int H, int PAIR_DIM> __global__ void fused_kernel(...) But since S=8, H=32, PAIR_DIM=64 are fixed by the problem's inputs? Actually they are constructor args (hidden_dim, pair_dim). They can vary? In the reference, msa_dim, pair_dim, hidden_dim are constructor arguments. The get_init_inputs returns [msa_dim, pair_dim, hidden_dim]. Those are the specific sizes used for evaluation. However, to be safe, we can make the kernel generic with runtime dimensions. But shared memory size for O_shared depends on H*H. H is hidden_dim = 32. That is a compile-time constant if we want to use static shared memory. We can use dynamic shared memory: extern __shared__ float shared[]; and partition it. But simpler: we can allocate enough shared memory for maximum expected H (maybe 32 or 64). Since hidden_dim=32 in this problem, we can allocate [8*64] for a and b and [64*64] for O. But to be generic, we can pass dimensions and use dynamic shared memory. However, dynamic shared memory in extern __shared__ float shared[]; must be declared in the kernel. Then we partition: float* a_shared = shared; float* b_shared = shared + S*H; float* o_shared = b_shared + S*H. But S and H are not compile-time. That's fine. Wait, but if we want to avoid dynamic shared memory, we can just put the arrays with max sizes. But better to use dynamic for flexibility. Alternatively, we can compute O on the fly without storing in shared memory: each thread computes a portion of O, but we need O for all d. Let's think about other approaches. Approach without O_shared: Each thread (d) loops over h1, h2. For each (h1,h2), compute sum_s = ... (redundant). This avoids needing O_shared and extra syncthreads. It's simpler. Weight access is coalesced. The redundant computation is 8192 mul-adds per thread (vs 1024 if using O_shared). 8192 vs 1024 is 8x more compute. But compute might be fine if memory is the bottleneck? Actually 8192 per thread is not huge. Each block with 64 threads does 524k ops. With 2048 blocks, total 1.07B ops. At 15 TFLOPS, that's 0.07 ms. At 1 TFLOPS (if memory bound), 1 ms. Likely still much faster than reference. So simplicity might be better. But to maximize performance and reduce redundant compute, we can use O_shared. Let's evaluate the dynamic shared memory approach. Kernel signature: __global__ void outer_prod_mean_out_proj_kernel( const float* __restrict__ a, const float* __restrict__ b, const float* __restrict__ denom, const float* __restrict__ weight_t, // shape (H*H, pair_dim) const float* __restrict__ bias, float* __restrict__ out, int B, int S, int N, int H, int pair_dim) { int b = blockIdx.x; int i = blockIdx.y; int j = blockIdx.z; int d = threadIdx.x; extern __shared__ float sbuf[]; float* a_shared = sbuf; float* b_shared = a_shared + S * H; float* o_shared = b_shared + S * H; int a_base = ((b * S) * N + i) * H; int b_base = ((b * S) * N + j) * H; int total = S * H; for (int idx = d; idx < total; idx += blockDim.x) { int s = idx / H; int h = idx % H; a_shared[idx] = a[a_base + s * N * H + h]; b_shared[idx] = b[b_base + s * N * H + h]; } __syncthreads(); // compute O int hh = H * H; for (int idx = d; idx < hh; idx += blockDim.x) { int h1 = idx / H; int h2 = idx % H; float sum = 0.0f; for (int s = 0; s < S; ++s) { sum += a_shared[s * H + h1] * b_shared[s * H + h2]; } float denom_val = denom[b * N + i]; // denom should be >= 1.f o_shared[idx] = sum / denom_val; } __syncthreads(); float acc = (bias != nullptr) ? bias[d] : 0.0f; for (int idx = 0; idx < hh; ++idx) { float w = weight_t[idx * pair_dim + d]; acc += w * o_shared[idx]; } int out_idx = ((b * N + i) * N + j) * pair_dim + d; out[out_idx] = acc; } Shared memory size = (2*S*H + H*H) * sizeof(float). With S=8, H=32: (512 + 1024)*4 = 6144 bytes. This seems solid. But wait: denom must be at least 1. We can compute denom_clamped = denom.clamp(min=1) in Python. Also need to consider: the reference uses a = a * mask and b = b * mask before einsum. So a and b passed to kernel must already be masked. We can do a = a * mask; b = b * mask in Python. That's fine. But we could also pass mask to kernel and apply there. However, the elementwise multiply is cheap and PyTorch does it efficiently. Let's keep it simple. Potential numerical issues: LayerNorm and linear projections are done in PyTorch as before. The custom kernel computes sum_s of a*b / denom. The reference computes einsum over s and then divides. In the reference, a and b are already masked. So einsum yields sum_s a[s]*b[s]. Then outer / denom. The kernel does the same. Good. But note: denom is sum over dim=1 of mask, unsqueezed to match. In reference they do denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1) before division. Then outer = outer / denom. Since denom is broadcast over j,h1,h2. In kernel we divide each O element by denom_i. Good. Let's build the Python code. ModelNew.__init__: - Initialize layer_norm, proj_a, proj_b, out_proj with same args. We can just copy weights from Model? No, we just create new nn.Module and parameters. The test will instantiate ModelNew with same args and likely load same random weights? Wait the evaluation may test correctness by copying weights? Usually they initialize ModelNew and then run with same weights? Actually the evaluation harness probably creates Model with random weights and gets outputs, then creates ModelNew with same init inputs and runs forward on same inputs. But ModelNew's parameters are randomly initialized, so outputs will differ. How do they check correctness? Perhaps they check architecture equivalence? But the prompt says "Replace the forward computation of the PyTorch model below with a custom CUDA kernel. Your ModelNew class must: 1. Accept the same constructor arguments as Model. 2. Implement forward() with the same signature. 3. Produce numerically equivalent outputs at fp32 precision." Typically in these benchmarks, they copy the weights from the reference model to ModelNew before evaluating. Or they evaluate the init and forward independently. How is it done? We need to ensure that ModelNew has the same parameters and that they are used correctly. Usually the framework does something like: model = Model(*get_init_inputs()) model_new = ModelNew(*get_init_inputs()) # copy weights model_new.load_state_dict(model.state_dict()) Or they initialize both with same random seed. The prompt says "The reference implementation" includes Model class. The evaluation likely instantiates Model and ModelNew and copies the state dict. So we need to have matching parameter names and shapes. Our ModelNew can use nn.LayerNorm, nn.Linear same as Model. That will create same parameter names. So state dict copying should work. ModelNew.forward: - x = self.layer_norm(msa_repr) - a = self.proj_a(x) - b = self.proj_b(x) - mask = msa_mask.unsqueeze(-1) - a = a * mask - b = b * mask - denom = msa_mask.sum(dim=1).clamp(min=1) # shape (B,N) - weight_t = self.out_proj.weight.t().contiguous() # shape (pair_dim, hidden_dim*hidden_dim) -> transposed to (hidden_dim*hidden_dim, pair_dim) - bias = self.out_proj.bias # shape (pair_dim,) - output = torch.empty(B, N, N, pair_dim, device=..., dtype=torch.float32) - launch kernel. Wait: out_proj weight shape is (pair_dim, H*H). After transpose, it's (H*H, pair_dim). Good. The kernel code uses `weight_t` (const float*). Now, let’s write the CUDA code. We need to use `torch.utils.cpp_extension.load_inline`. We need to provide `cpp_sources` and `cuda_sources`. `load_inline` auto-generates pybind11 module based on functions list. Example: import torch from torch.utils.cpp_extension import load_inline cuda_source = r''' __global__ void kernel(...) { ... } ''' cpp_source = r''' torch::Tensor forward(torch::Tensor a, torch::Tensor b, torch::Tensor denom, torch::Tensor weight_t, torch::Tensor bias, int pair_dim) { int B = a.size(0); int S = a.size(1); int N = a.size(2); int H = a.size(3); auto out = torch::empty({B, N, N, pair_dim}, a.options()); dim3 grid(B, N, N); dim3 block(pair_dim); size_t shared = (2 * S * H + H * H) * sizeof(float); kernel<<<grid, block, shared>>>( a.data_ptr<float>(), b.data_ptr<float>(), denom.data_ptr<float>(), weight_t.data_ptr<float>(), bias.defined() ? bias.data_ptr<float>() : nullptr, out.data_ptr<float>(), B, S, N, H, pair_dim ); return out; } ''' We need to ensure pair_dim matches bias size and block size. pair_dim is known at runtime (constructor arg). The block size must be at most 1024. pair_dim=64 is fine. But if pair_dim were larger (e.g., 128), block size would be 128, still okay. If pair_dim > 1024, we can't launch. But the problem's pair_dim is 64. We can assert pair_dim <= 1024. Alternatively, we can make the block size a constant (e.g., 64) and loop over d. But that adds complexity. Let's keep it flexible but assert. We need to handle cases where denom is on CPU? No, denom should be on same device as a. Let's write the code carefully. Important: denom shape is (B,N). We need to ensure it's contiguous and float. Denom = msa_mask.sum(dim=1).clamp(min=1).to(a.dtype).contiguous() Wait, msa_mask can be float or int? In get_inputs, msa_mask = torch.ones(...), default float32. sum yields float32. Good. But to be safe, do denom = msa_mask.sum(dim=1).float().clamp(min=1.0).contiguous() Now, bias might be None if out_proj.bias is False. But by default, nn.Linear has bias=True. We can handle it in C++ wrapper: pass bias tensor or optional. In pybind11, optional<torch::Tensor> can be used. But since load_inline auto-generates bindings from the function signature, we need to keep it simple. We can always pass bias tensor; if None, pass an empty tensor? Or we can handle in Python: bias = self.out_proj.bias if self.out_proj.bias is not None else torch.zeros(pair_dim, device=..., dtype=...) Simpler: In Python, if bias is None, create a zero tensor. Or just always pass self.out_proj.bias if it exists else a zero tensor. Let's just require bias exists. The reference uses default nn.Linear which has bias. So we are fine. Now, compile and test. One nuance: `load_inline` needs to compile. We'll write the code as a string in Python. Need to be careful about escaping. Let's draft the kernel code in the analysis and then compile. Potential issues: - The reference also divides by denom before reshape. Our kernel does that. - Numerical tolerance: fp32 atol=rtol=1e-4. Using float arithmetic as described should be fine. - The layer norm and linear layers: we should use the same PyTorch modules to avoid reimplementing them. That's fine. Let's write the code: ```python import torch import torch.nn as nn import math from torch.utils.cpp_extension import load_inline # Load the CUDA kernel cuda_source = r''' __global__ void outer_prod_mean_out_proj_kernel( const float* __restrict__ a, const float* __restrict__ b, const float* __restrict__ denom, const float* __restrict__ weight_t, // shape (H*H, pair_dim) const float* __restrict__ bias, float* __restrict__ out, int B, int S, int N, int H, int pair_dim) { int b = blockIdx.x; int i = blockIdx.y; int j = blockIdx.z; int d = threadIdx.x; // Dynamic shared memory layout: a_shared[S*H], b_shared[S*H], o_shared[H*H] extern __shared__ float sbuf[]; float* a_shared = sbuf; float* b_shared = a_shared + S * H; float* o_shared = b_shared + S * H; int a_base = ((b * S) * N + i) * H; int b_base = ((b * S) * N + j) * H; int total = S * H; for (int idx = d; idx < total; idx += blockDim.x) { int s = idx / H; int h = idx % H; a_shared[idx] = a[a_base + s * N * H + h]; b_shared[idx] = b[b_base + s * N * H + h]; } __syncthreads(); float denom_val = denom[b * N + i]; // denom_val is already clamped >= 1 on host side int hh = H * H; for (int idx = d; idx < hh; idx += blockDim.x) { int h1 = idx / H; int h2 = idx % H; float sum = 0.0f; #pragma unroll for (int s = 0; s < S; ++s) { sum += a_shared[s * H + h1] * b_shared[s * H + h2]; } o_shared[idx] = sum / denom_val; } __syncthreads(); float acc = (bias != nullptr) ? bias[d] : 0.0f; for (int idx = 0; idx < hh; ++idx) { float w = weight_t[idx * pair_dim + d]; acc += w * o_shared[idx]; } int out_idx = ((b * N + i) * N + j) * pair_dim + d; out[out_idx] = acc; } ''' cpp_source = r''' torch::Tensor outer_prod_mean_out_proj( torch::Tensor a, torch::Tensor b, torch::Tensor denom, torch::Tensor weight_t, torch::Tensor bias, int pair_dim) { int B = a.size(0); int S = a.size(1); int N = a.size(2); int H = a.size(3); auto out = torch::empty({B, N, N, pair_dim}, a.options()); dim3 grid(B, N, N); dim3 block(pair_dim); size_t shared_size = (2 * S * H + H * H) * sizeof(float); outer_prod_mean_out_proj_kernel<<<grid, block, shared_size>>>( a.data_ptr<float>(), b.data_ptr<float>(), denom.data_ptr<float>(), weight_t.data_ptr<float>(), bias.defined() ? bias.data_ptr<float>() : nullptr, out.data_ptr<float>(), B, S, N, H, pair_dim ); return out; } ''' The CUDA code is compiled with nvcc. #pragma unroll is fine. In the cpp wrapper, we need to ensure `weight_t` and `bias` are contiguous float tensors. We can do that in Python. Wait, there is also `denom`. In the kernel, denom is a float pointer. In Python we can ensure it's contiguous float. But there is a subtlety: the denom shape is (B,N). In the kernel, we read denom[b * N + i]. That's fine if contiguous. Now, ModelNew class: ```python class ModelNew(nn.Module): def __init__(self, msa_dim, pair_dim, hidden_dim=32): super().__init__() self.layer_norm = nn.LayerNorm(msa_dim) self.proj_a = nn.Linear(msa_dim, hidden_dim) self.proj_b = nn.Linear(msa_dim, hidden_dim) self.out_proj = nn.Linear(hidden_dim * hidden_dim, pair_dim) def forward(self, msa_repr, msa_mask): x = self.layer_norm(msa_repr) a = self.proj_a(x) b = self.proj_b(x) mask = msa_mask.unsqueeze(-1) a = a * mask b = b * mask denom = msa_mask.sum(dim=1).clamp(min=1.0).to(a.dtype).contiguous() weight_t = self.out_proj.weight.t().contiguous() bias = self.out_proj.bias.contiguous() if self.out_proj.bias is not None else torch.zeros(self.out_proj.out_features, device=a.device, dtype=a.dtype) # Ensure all inputs are contiguous and float a = a.contiguous() b = b.contiguous() out = outer_prod_mean_out_proj(a, b, denom, weight_t, bias, self.out_proj.out_features) return out ``` Wait, `self.out_proj.out_features` is pair_dim. That matches. One issue: `outer_prod_mean_out_proj` is defined in the loaded extension module. We need to assign it to a Python function. Code: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline # compile _fused_op = load_inline( name="outer_prod_mean", cpp_sources=cpp_source, cuda_sources=cuda_source, functions=["outer_prod_mean_out_proj"], extra_cuda_cflags=["--use_fast_math"], verbose=False, ) ``` Then call `_fused_op.outer_prod_mean_out_proj(...)`. But we need to be careful: `load_inline` caches the compiled extension by name. If we change the code and rerun, it may use cache. For submission, that's fine. Let's test compile.