Now let me analyze the problem: - Input: h0 of shape (16, 48) - batch size 16, input features 48 - W1: (48, 64) -> h1: (16, 64) with ReLU - W2: (64, 40) -> h2: (16, 40) with ReLU - W3: (40, 24) -> out: (16, 24) This is a small MLP. The batch size is only 16, so there's limited parallelism across the batch. Let me calculate the FLOPs: - Layer 1: 16 * 48 * 64 * 2 = 98,304 FLOPs - Layer 2: 16 * 64 * 40 * 2 = 81,920 FLOPs - Layer 3: 16 * 40 * 24 * 2 = 30,720 FLOPs - Total: ~211K FLOPs This is tiny! The latency overhead might dominate. Let me write a fused kernel. Given the small sizes, a single kernel that does all three layers with shared memory might be optimal. Each thread block can handle one batch element. Actually, with batch size 16 and small dimensions, let me think about the best approach: 1. Option 1: Single fused kernel where each block handles one sample - 16 blocks, each with enough threads to parallelize the matrix multiplies - Can use shared memory to avoid global memory round trips 2. Option 2: Three separate kernels (like PyTorch) - More kernel launch overhead For small sizes, the fused kernel should win because it reduces kernel launch overhead and global memory traffic. Let me write a fused kernel where each thread block processes one batch element: - Block 0 processes batch element 0 - Block 1 processes batch element 1 - etc. Within each block: - We need to compute h0[batch_idx] @ W1 -> h1 (48x64 multiply) - Then ReLU h1 - Then h1 @ W2 -> h2 (64x40 multiply) - Then ReLU h2 - Then h2 @ W3 -> out (40x24 multiply) For the matrix multiplies, we can use a tiled approach or just have each thread compute one output element. Given the small sizes: - H0=48, H1=64, H2=40, H3=24 - B=16 For the first layer (48x64), we have 64 output elements per batch. We could use 64 threads per block to compute one output each. For the second layer (64x40), we have 40 output elements per batch. For the third layer (40x24), we have 24 output elements per batch. Actually, let me think about this differently. The maximum threads per block should handle the largest layer output. That's 64 for the first layer. So let's use 64 threads per block (one warp + more), and have each thread compute one output element of each layer sequentially. For layer 1: Each thread computes h1[j] = sum over k of h0[k] * W1[k,j], then applies ReLU. For layer 2: Each thread (only first 40 are active) computes h2[j] = sum over k of h1[k] * W2[k,j], then ReLU. For layer 3: Each thread (only first 24 are active) computes out[j] = sum over k of h2[k] * W3[k,j]. Wait, the problem is that we need all 64 values of h1 before we can compute h2. So there's a synchronization point. Let me use shared memory to store intermediate results: 1. Load h0[batch_idx] into shared memory (48 floats) - could be in registers actually 2. Each thread j computes h1[j] and stores in shared memory (64 floats) 3. __syncthreads() 4. Each thread j (j < 40) computes h2[j] and stores in shared memory (40 floats) 5. __syncthreads() 6. Each thread j (j < 24) computes out[j] and writes to global memory This should work! Let me code it up. Actually, let me reconsider the thread count. We need at least 64 threads for the first layer, and we need to handle all three layers. Using 64 threads per block (2 warps) should be efficient. Actually, wait. With B=16 batch elements and small matrices, the total data is tiny. Let me think about whether we should process multiple batch elements per block. With the current design (1 batch per block), we'd have 16 blocks running. That's not much parallelism but might be okay on H100. Alternatively, we could process all 16 batch elements in a single block with many threads. But that might be complex. Let me start with the simpler design: one block per batch element. Here's my plan: - Grid: dim3(B, 1, 1) = dim3(16, 1, 1) - Block: dim3(64, 1, 1) = dim3(64, 1, 1) Shared memory per block: - h0_local: 48 floats = 192 bytes (could use registers) - h1_local: 64 floats = 256 bytes - h2_local: 40 floats = 160 bytes - Total: ~608 bytes (well within limits) Let me write this kernel: ```cpp __global__ void mlp_fused_kernel( const float* __restrict__ h0, // Input: (B, H0) const float* __restrict__ W1, // Weight: (H0, H1) const float* __restrict__ W2, // Weight: (H1, H2) const float* __restrict__ W3, // Weight: (H2, H3) float* __restrict__ out, // Output: (B, H3) int B, int H0, int H1, int H2, int H3 ) { int batch_idx = blockIdx.x; int thread_idx = threadIdx.x; // Shared memory for intermediate results __shared__ float h0_local[48]; // H0 __shared__ float h1_local[64]; // H1 __shared__ float h2_local[40]; // H2 // Cooperative load of h0 for this batch if (thread_idx < H0) { h0_local[thread_idx] = h0[batch_idx * H0 + thread_idx]; } __syncthreads(); // Layer 1: h1 = ReLU(h0 @ W1) // Each thread computes one element of h1 if (thread_idx < H1) { float sum = 0.0f; for (int k = 0; k < H0; k++) { sum += h0_local[k] * W1[k * H1 + thread_idx]; } h1_local[thread_idx] = fmaxf(sum, 0.0f); } __syncthreads(); // Layer 2: h2 = ReLU(h1 @ W2) // Each thread computes one element of h2 if (thread_idx < H2) { float sum = 0.0f; for (int k = 0; k < H1; k++) { sum += h1_local[k] * W2[k * H2 + thread_idx]; } h2_local[thread_idx] = fmaxf(sum, 0.0f); } __syncthreads(); // Layer 3: out = h2 @ W3 // Each thread computes one element of out if (thread_idx < H3) { float sum = 0.0f; for (int k = 0; k < H2; k++) { sum += h2_local[k] * W3[k * H3 + thread_idx]; } out[batch_idx * H3 + thread_idx] = sum; } } ``` Wait, I need to double-check the weight layout. In PyTorch, nn.Linear(H0, H1) creates a weight of shape (H1, H0), and the operation is `input @ weight.T` which is equivalent to `input @ weight.T`. So h0 @ W1.T gives h1 of shape (B, H1), where W1 is (H1, H0). Let me recalculate: - h0 is (B, H0) = (16, 48) - W1 is (H1, H0) = (64, 48) - h1 = h0 @ W1.T = (16, 48) @ (48, 64) = (16, 64) - W2 is (H2, H1) = (40, 64) - h2 = h1 @ W2.T = (16, 64) @ (64, 40) = (16, 40) - W3 is (H3, H2) = (24, 40) - out = h2 @ W3.T = (16, 40) @ (40, 24) = (16, 24) So W1 has shape (64, 48), stored in row-major as W1[i, j] = W1[i * 48 + j] for output i and input j. For computing h1[b, j] = sum_k h0[b, k] * W1[j, k] = sum_k h0[b, k] * W1[j * H0 + k] This is different from what I wrote. Let me fix it: ```cpp // For h1[j] (j is output index): // h1[j] = sum over k of h0[k] * W1[j, k] = sum over k of h0[k] * W1[j * H0 + k] ``` So the kernel should be: ```cpp __global__ void mlp_fused_kernel( const float* __restrict__ h0, // Input: (B, H0) const float* __restrict__ W1, // Weight: (H1, H0) const float* __restrict__ W2, // Weight: (H2, H1) const float* __restrict__ W3, // Weight: (H3, H2) float* __restrict__ out, // Output: (B, H3) int B, int H0, int H1, int H2, int H3 ) { int batch_idx = blockIdx.x; int thread_idx = threadIdx.x; // Shared memory for intermediate results __shared__ float h0_local[48]; // H0 __shared__ float h1_local[64]; // H1 __shared__ float h2_local[40]; // H2 // Cooperative load of h0 for this batch if (thread_idx < H0) { h0_local[thread_idx] = h0[batch_idx * H0 + thread_idx]; } __syncthreads(); // Layer 1: h1 = ReLU(h0 @ W1.T) // Each thread computes one element of h1 if (thread_idx < H1) { float sum = 0.0f; for (int k = 0; k < H0; k++) { // W1 has shape (H1, H0), W1[thread_idx, k] = W1[thread_idx * H0 + k] sum += h0_local[k] * W1[thread_idx * H0 + k]; } h1_local[thread_idx] = fmaxf(sum, 0.0f); } __syncthreads(); // Layer 2: h2 = ReLU(h1 @ W2.T) // Each thread computes one element of h2 if (thread_idx < H2) { float sum = 0.0f; for (int k = 0; k < H1; k++) { // W2 has shape (H2, H1), W2[thread_idx, k] = W2[thread_idx * H1 + k] sum += h1_local[k] * W2[thread_idx * H1 + k]; } h2_local[thread_idx] = fmaxf(sum, 0.0f); } __syncthreads(); // Layer 3: out = h2 @ W3.T // Each thread computes one element of out if (thread_idx < H3) { float sum = 0.0f; for (int k = 0; k < H2; k++) { // W3 has shape (H3, H2), W3[thread_idx, k] = W3[thread_idx * H2 + k] sum += h2_local[k] * W3[thread_idx * H2 + k]; } out[batch_idx * H3 + thread_idx] = sum; } } ``` This looks correct now. Let me write the complete Python wrapper.