Let me analyze the reference model: 1. `self.gemm = nn.Linear(in_features, out_features)` - matrix multiplication: [16384, 4096] @ [4096, 4096] + [4096] 2. `x = x * self.scale` - element-wise multiply by scale [4096] 3. `self.bn = nn.BatchNorm1d(out_features, eps=eps, momentum=0.1)` - batch norm along dimension 0 for each feature Wait, actually `nn.BatchNorm1d(out_features)` normalizes across the batch dimension (dim=0) for each feature. The operation is: - y = x @ W.T + b - y = y * scale - y = BN(y) = (y - mean) / sqrt(var + eps) * gamma + beta Where gamma and beta are the BN parameters, and mean/var are running stats during eval. But wait - the reference doesn't set eval mode, so in training mode, BN computes mean and var from the current batch. Actually, looking more carefully: the default mode is training. So for each forward pass: 1. Compute z = x @ W^T + b 2. Compute y = z * scale 3. Compute batch mean and var of y 4. Normalize y and apply BN's gamma/beta 5. Update running stats Wait, the problem says the model is run in default mode. But the evaluation typically runs in some mode. Let me think... For correctness testing, the reference model will be called with the same inputs. I need to handle whatever mode it's in. Actually, looking at the default behavior, if the model is in training mode, BN computes statistics per-batch. But typically for custom kernels, they test with the same model state. Let me approach this by fusing as many operations as possible into a single kernel. Actually, a key insight: we can fuse the linear + scale + batch norm into a single kernel when in training mode, or into a fused kernel in general. But wait - in training mode, BN uses batch statistics. Let me re-read the reference... The reference doesn't call .train() or .eval(), so it defaults to training mode. But in evaluation, the kernel tests might run the model in a specific mode. Actually, for the benchmarking and correctness, the model state (weights, running_mean, running_var, training/eval mode) is whatever it gets initialized to. By default, nn.Module is in training mode. So for correctness, I need to match PyTorch's training mode exactly: 1. Compute output of linear: out = x @ W^T + b 2. Scale: out = out * scale 3. BN training mode: - Compute batch mean: mean = out.mean(dim=0) - Compute batch var: var = out.var(dim=0, unbiased=False) or unbiased=True? Wait, PyTorch BatchNorm1d during training uses unbiased=False for var computation (Bessel's correction not applied). Let me verify. Actually, PyTorch's BatchNorm uses `var = input.var(dim, unbiased=False)` during training for the normalization denominator, but the running_var is updated with unbiased variance. Let me be more precise. PyTorch source: - During training, it computes mean and var over N, H, W (for NCHW) - mean = x.mean([0, 2, 3]) - var = x.var([0, 2, 3], unbiased=False) - Then normalizes: (x - mean) / sqrt(var + eps) - Updates running_mean = (1 - momentum) * running_mean + momentum * mean - Updates running_var = (1 - momentum) * running_var + momentum * var * n / (n - 1) [unbiased correction] Wait no, let me check again. Actually I believe: - training var for normalization: unbiased=False - running_var update: uses unbiased variance, i.e., var * n / (n-1) So we need to compute: 1. out = matmul(x, W.T) + b 2. out = out * scale 3. mean = out.mean(dim=0) 4. var = out.var(dim=0, unbiased=False) 5. out = (out - mean) / sqrt(var + eps) * gamma + beta 6. Update running_mean, running_var For a naive fused kernel, I can: 1. First compute the linear result into a temporary 2. Apply scaling 3. Compute mean and var (can be done in shared memory reduction) 4. Normalize But with matrices of size [16384, 4096], this is: - Matmul: 16384 * 4096 * 4096 * 2 = ~549 GFLOPs - BN: element-wise ops, very small The matmul dominates. So the best approach is probably: 1. Use cuBLAS for the matmul (via PyTorch's linear) 2. Fuse scale + BN into a custom kernel Or we can try to use a custom CUTLASS-like kernel, but on H100, cuBLAS is extremely optimized. Actually, we might be able to fuse scale into the bias addition, and then do BN. But the scale multiplies the output, not just the bias. y = (x @ W^T + b) * scale = x @ W^T * scale + b * scale = x @ (W * scale)^T + b * scale So we can fuse the scale into the linear layer weights! That's a big optimization because we avoid an extra memory pass. Then we just need to apply BN to the result. So: 1. Fuse scale into W and b 2. Call linear with fused weights 3. Apply BN in a custom kernel (computing mean, var, normalizing) Wait, but we can't modify the original weights permanently because the model should maintain its parameters. But we can precompute fused weights in __init__ or lazily in forward. Actually, the reference model stores W and b separately from scale. If I fuse them in forward, I'd need to compute W_fused = W * scale and b_fused = b * scale each time. For [4096, 4096] weights, that's 16M elements, which is not too bad. But it adds a memory read/write. Alternatively, I could do the matmul with original W, then in my BN kernel, first multiply by scale, then compute mean and variance, then normalize. But that requires reading and writing the entire output again. Wait, let me think more carefully: Option A: Fuse scale into linear - Create W_fused = W * scale (read W, read scale, write W_fused: ~48MB memory traffic) - Create b_fused = b * scale (read b, read scale, write b_fused: negligible) - Call linear with x, W_fused, b_fused - Apply BN on result Option B: Custom kernel for scale + BN - Call linear normally - In custom kernel: load elements, multiply by scale, compute mean/var (requires multiple passes or clever reduction), then normalize For Option B with a single kernel: - We need to compute mean and variance first, which requires either two passes or a persistent kernel approach. - For a tensor of size [16384, 4096], we can use a two-pass approach in CUDA: - Kernel 1: Compute mean (reduction) - Kernel 2: Compute var using mean (reduction) - Kernel 3: Normalize Or we can use warp shuffle / shared memory reductions. But actually, PyTorch's native batch norm is already highly optimized (it uses cuDNN which fuses everything). So why would a custom kernel be faster? Looking at the reference: ```python x = self.gemm(x) # linear, uses cuBLAS x = x * self.scale # element-wise, memory bound x = self.bn(x) # cuDNN fused BN, but needs separate memory for input ``` The overheads are: 1. Writing output of linear to global memory 2. Reading/writing for scale multiplication 3. BN reads the scaled output again If we can: 1. Keep linear output in registers/shared memory and fuse with scale 2. Then do BN all in one go But linear outputs one element per thread (at best). Actually, for a matmul of [16384, 4096] x [4096, 4096], we can compute the output with multiple warps per row or whatever. Actually, let me reconsider. The matmul is huge - 549 GFLOPs. At 67 TFLOPS FP32, that's ~8.2ms. Memory bandwidth is 3.35 TB/s, so reading the 256MB output for scale+BN at memory bandwidth takes ~76 microseconds. So the matmul dominates. Fusing scale into the linear weights or into the linear computation won't save that much relative to the matmul time. But wait - cuBLAS on H100 with tensor cores in TF32 mode can do ~989 TFLOPS. The matmul in TF32 would take ~0.55ms. At that speed, the memory-bound operations matter more. Actually, PyTorch's nn.Linear in fp32 might not use TF32 by default. Actually, it depends on the settings. But for fp32 precision requirements, we need to be careful. Wait, the correctness tolerance is atol=rtol=1e-4 for fp32. TF32 has ~1e-3 precision typically. So we might not be able to use TF32 if we want to meet fp32 tolerance. Actually, PyTorch's matmul will use TF32 when `torch.backends.cuda.matmul.allow_tf32` is True (default in newer PyTorch). But the reference model will do the same thing, so as long as we use the same computation, it's fine. Hmm, but if we're writing custom CUDA, we can just use cuBLAS via torch.mm or similar. Actually, I think the key insight is that the reference does: 1. Linear matmul 2. Scale (element-wise multiply) 3. BatchNorm (computes stats, normalizes) For PyTorch's cuDNN BN, it probably fuses the stats computation and normalization. But scale is separate. If we fuse scale + BN, we eliminate one memory pass. But actually, can we do better? Let me think about the entire thing. For training mode BN on [16384, 4096]: - Batch size = 16384, features = 4096 - Computing mean over 16384 elements per feature - Computing var over 16384 elements per feature - Normalizing all elements With 4096 features, each feature has 16384 elements. We can parallelize across features. Let me write a kernel that does: 1. Use torch.matmul for the linear part 2. Fused kernel for "multiply by scale + compute mean + compute var + normalize" For the fused kernel: - We need to compute mean and var first, then normalize. Since it's a single kernel launch, we need to do this in multiple stages using shared memory and __syncthreads(). Actually, for coalesced access, if the output is [N, C] = [16384, 4096] in row-major, then elements of the same feature are strided by 4096. That's not very coalesced if each thread reads one element per feature. Wait, if the layout is [N, C], then for a fixed feature c, the elements are at positions c, C+c, 2C+c, etc. So adjacent threads accessing the same feature are strided by C=4096. For fp32, that's 16KB apart. That's not coalesced. For reduction over N with fixed C, it's better to have each thread process adjacent elements (so threads in a warp access elements at c, c+1, c+2, etc. but that means each thread is processing a different feature). Actually, for computing mean across N for each feature: - If we assign one block per feature, then threads in the block read strided elements. - For coalesced access, we should consecutive threads read consecutive memory. - So thread i should read element (some_n * C + some_c), and adjacent threads should read adjacent elements in memory. - This means adjacent threads should work on the same row but adjacent columns. So for a reduction over rows (computing mean per column): - Each block handles one or more columns. - Threads read consecutive columns in the same row, then move to the next row. - Within a block, we compute the sum for each column using shared memory. This is much more memory-friendly. Let me design the kernel: ```cuda template <typename T> __global__ void scale_bn_forward_kernel( const T* input, // [N, C] const T* scale, // [C] const T* gamma, // [C] (bn weight) const T* beta, // [C] (bn bias) T* output, // [N, C] T* mean, // [C] T* var, // [C] T* running_mean, T* running_var, int N, int C, float eps, float momentum ) { // Each block handles one or more columns // Let's have each block handle BLOCK_COLS columns // For maximum simplicity, each block handles one column int c = blockIdx.x; if (c >= C) return; // Compute mean for column c float sum = 0.0f; for (int i = threadIdx.x; i < N; i += blockDim.x) { float val = input[i * C + c] * scale[c]; sum += val; } // Warp reduction then block reduction __shared__ float s_mean; // ... reduction ... if (threadIdx.x == 0) { s_mean = sum / N; mean[c] = s_mean; } __syncthreads(); // Compute var float sq_diff = 0.0f; for (int i = threadIdx.x; i < N; i += blockDim.x) { float val = input[i * C + c] * scale[c]; float diff = val - s_mean; sq_diff += diff * diff; } // ... reduction ... if (threadIdx.x == 0) { float variance = sq_diff / N; var[c] = variance; // Normalize and write output float inv_std = rsqrtf(variance + eps); // Update running stats running_mean[c] = (1 - momentum) * running_mean[c] + momentum * s_mean; running_var[c] = (1 - momentum) * running_var[c] + momentum * variance * N / (N - 1); } __syncthreads(); // Write normalized output float inv_std = rsqrtf(var[c] + eps); for (int i = threadIdx.x; i < N; i += blockDim.x) { float val = input[i * C + c] * scale[c]; output[i * C + c] = (val - mean[c]) * inv_std * gamma[c] + beta[c]; } } ``` Wait, there are several issues: 1. The reduction code is incomplete 2. The shared memory var[c] access after __syncthreads is racy 3. We compute the scaled value multiple times Let me refine. Actually, using one block per column with 1024 threads: - N = 16384, each thread handles 16 elements - Threads access memory with stride C = 4096. Adjacent threads are 4096*4 = 16384 bytes apart, which is not coalesced. For coalesced access, threads in a warp should access consecutive elements. But for a column reduction, consecutive elements are C apart. Hmm, but with 4096 columns, if we access strided by 4096, and each thread processes elements spaced by blockDim.x * C... Actually, on an H100 with high L2 cache bandwidth, strided access with 4096-element stride might still be okay because the cache line is 128 bytes = 32 floats. With 32 threads in a warp accessing elements spaced by 4096*4 = 16384 bytes, each thread touches a different cache line. That's 32 cache lines per warp = 4096 bytes, which is fine as long as the L2 can handle it. But it's definitely not ideal. Alternative: transpose the layout or process the data differently. Actually, for PyTorch tensors of shape [N, C], consecutive elements in a row are consecutive in memory. So element [n, c] is at offset n*C + c. For a column reduction, adjacent threads should ideally process the same row, adjacent columns. But then each thread is reducing a different column. Here's a better layout: - Block size: maybe 256 threads - Each block handles a tile of columns, say 32 columns - Thread i within the block processes row r, column c where c is in the tile - When reading, thread i reads [r, c0 + i%32] where 32 is consecutive columns - This is coalesced! Then within the block, we can use warp shuffles or shared memory to reduce across rows for each column. Actually, let's think about it more carefully: - We have 4096 columns. If each block handles 32 columns, we need 128 blocks. - Each block has 256 threads. - For reading row r, threads 0-31 read columns 0-31 (coalesced), threads 32-63 might read columns 0-31 of the next row or columns 32-63, etc. Actually, a more standard approach for column reduction on row-major matrices: - Launch a 2D grid of blocks - Each block computes partial sums for some columns - Then a second kernel reduces the partial sums Or for exactly one pass: - Grid: num_col_tiles x 1 - Block: some_threads - Each thread iterates over rows, reading consecutive columns in a coalesced manner. Let's say each block handles TILE_COLS = 32 columns. BlockDim = 256 threads. Each thread handles elements from different rows. For a fixed row r, we need all 32 columns read coalesced. So for row r, threads 0-31 read columns 0-31. For row r+1, threads 0-31 read columns 0-31 again? No, that gives 32 threads the same work. Actually, let's divide: 256 threads / 32 columns per block = 8 threads per column. So for column 0, threads 0, 32, 64, ..., 224 process it. For column 1, threads 1, 33, 65, ..., 225 process it. ... For column 31, threads 31, 63, ..., 255 process it. When reading row r: - Thread 0 reads [r, 0], thread 1 reads [r, 1], ..., thread 31 reads [r, 31]. - This is fully coalesced (consecutive memory). - Threads 32-63 read [r, 0] to [r, 31] again? No. Wait, each thread needs to process multiple rows. If thread t handles column c = t % 32, then for row r, thread t reads [r, c]. Adjacent threads read adjacent columns - perfect coalescing. Each thread processes rows r = (t / 32), r = (t / 32) + 8, r = (t / 32) + 16, etc. (since there are 8 threads per column). So thread t reads [r_start + k*8, c] where r_start = t / 32 and c = t % 32. For a warp of 32 threads: - Threads 0-31 read [r, 0] to [r, 31] for some r - coalesced! - Threads 32-63 - wait, that's the next warp. Actually, in the same warp (32 threads), all threads read the same row r = some value, and adjacent columns - perfectly coalesced. Then next iteration, they read row r+8 again. Also coalesced. This is great! Let me write a cleaner design: ```cuda #define TILE_COLS 32 __global__ void bn_forward_kernel( const float* input, // [N, C] const float* scale, const float* gamma, const float* beta, float* output, float* mean, float* var, float* running_mean, float* running_var, int N, int C, float eps, float momentum ) { const int tid = threadIdx.x; const int warps_per_block = blockDim.x / 32; const int warp_id = tid / 32; const int lane_id = tid % 32; const int col_block = blockIdx.x; const int c_start = col_block * TILE_COLS; // Each warp handles a subset of rows for all TILE_COLS columns // We'll have each thread handle elements spaced by warps_per_block in rows // Step 1: Compute partial sums for mean float local_sum[TILE_COLS] = {0}; for (int r = warp_id; r < N; r += warps_per_block) { int base_idx = r * C + c_start; #pragma unroll for (int c = 0; c < TILE_COLS && c_start + c < C; c++) { float val = input[base_idx + c] * scale[c_start + c]; local_sum[c] += val; } } // Step 2: Reduce local_sum across warp using shuffle __shared__ float s_sum[TILE_COLS]; // Actually need shared memory per column // ... } ``` Hmm, this is getting complex. Let me simplify. Actually, for N=16384 and C=4096, using one block per column with 512 threads and accepting strided access might still be fast enough, especially on H100 with its 3.35 TB/s memory bandwidth. Let me calculate: each element is 4 bytes. For computing mean + var + normalize, we read the input once (to compute mean), read it again (to compute var), and write the output once. With scale multiply fused in: - Read input: N*C * 4 bytes - Read scale: C * 4 bytes (negligible) - Read gamma, beta: C * 4 bytes (negligible) - Write output: N*C * 4 bytes Total: ~3 * 256MB = 768MB. At 3.35 TB/s, that's ~230 microseconds. For PyTorch reference: - Linear: uses tensor cores, but let's say ~1ms - Scale: read + write = 512MB, ~153 microseconds - BN: internally reads, computes, writes. Probably ~200-300 microseconds So the memory ops after linear are a small fraction. The key is to fuse them to reduce overhead. Actually, looking at this differently: can we fuse everything including the matmul? For a matmul of [16384, 4096] x [4096, 4096], we could write a custom kernel that: 1. Computes partial dot products 2. Multiplies by scale 3. Applies BN normalization But writing a high-performance matmul for 4096x4096 on H100 is extremely difficult to beat cuBLAS. However, we can potentially do a custom kernel that computes the matmul and then immediately does some processing, but cuBLAS expects to write to global memory. Wait, actually there's another approach. cuBLAS can write to global memory. Then in our fused kernel we read it back to do scale + BN. But we can at least avoid the separate scale kernel. Actually, on H100 with 80GB of memory and huge bandwidth, the overhead of multiple kernels isn't just memory - it's also kernel launch overhead and not fusing operations that prevent good instruction pipelining. Let me think about a different approach: can we use the fact that BN is an affine transform to fold it into the linear layer? The full computation is: out = BN((x @ W.T + b) * scale) = [((x @ W.T + b) * scale - mean) / sqrt(var + eps)] * gamma + beta During training, mean and var depend on the input x, so we can't precompute them. But during inference: mean = running_mean, var = running_var out = (x @ W.T + b - running_mean/scale) * scale / sqrt(var + eps) * gamma + beta Wait, that's not quite right. Let me expand: (x @ W.T + b) * scale = x @ W.T * scale + b * scale Then BN normalize and scale: = [(x @ W.T * scale + b * scale) - mean] / sqrt(var + eps) * gamma + beta During inference (mean = running_mean, var = running_var): = [(x @ W.T * scale + b * scale) - running_mean] / sqrt(running_var + eps) * gamma + beta = x @ W.T * scale / sqrt(running_var + eps) * gamma + (b * scale - running_mean) / sqrt(running_var + eps) * gamma + beta So we can fold scale / sqrt(running_var + eps) * gamma into W, and fold (b * scale - running_mean) / sqrt(running_var + eps) * gamma + beta into b. But during training, mean and var are computed per-batch, so we can't fully fold them. We still need to compute stats. Let me reconsider the whole approach. Maybe the model is run in eval mode? Let me check the reference again - it's not specified, but by default nn.Module is in training mode. Wait, I should check what PyTorch's batch norm does during evaluation vs training. The problem doesn't specify a mode, so I need to support the same behavior as the reference. In the reference code, the model is instantiated and used as-is, which means it's in training mode. But looking more carefully at typical benchmark setups, they might run in eval mode to avoid training-specific randomness. However, the problem statement says "produces numerically equivalent outputs" and the reference doesn't call .eval(). I think I should write a kernel that works for both training and eval mode, matching PyTorch exactly. But actually, looking at the reference, the default nn.BatchNorm1d is created and the model is in train mode by default. For the correctness test, the model might be called in whatever state it's in. Let me handle both modes. Actually, let me look at the problem from a practical standpoint. The most likely evaluation scenario: they create Model(), load weights somehow (or init), call forward. Since .train() is default, it's in training mode. But wait, for correctness testing they might create the model and run a single forward pass. In training mode, BN computes batch statistics. For eval mode, it uses running stats. Actually, I realize now - I should write my kernel to handle both. In my __init__, I can save the training mode and replicate PyTorch's behavior exactly. Let me just support both modes in the kernel: - If training: compute batch mean/var, update running stats, normalize - If eval: use running mean/var, normalize Now, for the actual kernel implementation, let me write a clean, efficient fused kernel. Actually, you know what - let me just try a straightforward approach first: fuse scale + BN into one kernel, and use a simple column-parallel reduction algorithm. If that's fast enough, great. If not, optimize further. For the reduction, using one warp per column would be ideal because warp shuffle can reduce without shared memory sync. With 4096 columns and 32 threads per warp, we need 4096 warps = 128 blocks of 1024 threads (32 warps each). Each warp handles one column. Wait, but 128 blocks might not fully utilize the GPU. H100 has 132 SMs or so. So 128 blocks might underutilize a bit, but it's close. Each warp = 32 threads, handles 1 column, processes N=16384 elements. Each thread processes 16384/32 = 512 elements. Memory access: thread i in warp c reads elements at [i, c], [i+32, c], [i+64, c], etc. Stride between consecutive threads: C * sizeof(float) = 16384 bytes = 16KB. On H100, warps can handle this strided access, but it's not optimal. Alternatively: use more threads per column but do block-level reduction. With 1024 threads per block and each block handling multiple columns. Actually, for better memory coalescing, we can transpose the data or use a tiled approach. But for the first version, let's see if strided is fast enough. Wait, I just realized: if we transpose the linear output first, then the reduction would be coalesced. But transposing costs O(N*C) memory operations. Let me calculate if transposing is worth it: - Transpose: read 256MB, write 256MB = 512MB traffic - Then reduction on transposed data (rows become columns) is coalesced: ~768MB traffic - Total: ~1280MB - Without transpose, strided access: ~768MB but with worse efficiency On H100, even at 50% bandwidth efficiency for strided, it's 1.675 TB/s effective. 768MB / 1.675 TB/s = 458 microseconds. That's probably fine. Actually, let me just try a simpler approach with shared memory reduction and see how fast it is. Here's my plan: 1. In ModelNew.__init__, store all parameters and buffers from the reference 2. In ModelNew.forward: - Call torch.matmul or F.linear for the matrix multiplication - Call custom fused kernel for "scale + BN" Wait, actually I can be smarter. Let me fuse the scale into the linear weight and bias! This is a valid transformation since: (x @ W.T + b) * scale = x @ (W * scale).T + b * scale So: - Precompute W_fused = W * scale, b_fused = b * scale - Call F.linear(x, W_fused, b_fused) - Then apply BN on the result This saves one memory pass! We just need to update W_fused and b_fused whenever the parameters change. But since this is inference (not training the model parameters), we can precompute in __init__ and update when needed. Wait, but the parameters might change between forward calls in the tests. Actually, for a benchmark, the parameters usually stay fixed. But to be safe, we can compute the fused weight on each forward pass - it's just element-wise multiply on 16M+4K elements, which is negligible. Actually, there's an even better approach: PyTorch's linear supports doing this via broadcasting: out = F.linear(x, weight * scale.unsqueeze(1), bias * scale) This computes W_fused and b_fused on the fly, then does the matmul. PyTorch might even fuse this into the matmul via some graph optimization, but even if not, the element-wise ops are tiny. Actually, let me check: `weight * scale.unsqueeze(1)` creates a temporary tensor of the same size as weight (16M elements). That's a memory allocation. But it's still tiny overhead. Alternatively, we can use `torch.addmm` with a custom bias, but with scaled weights we'd still need the temporary. Hmm, actually we can do the scale multiply in our custom BN kernel anyway, so pre-fusing into linear might not save much compared to fusing scale+BN. Wait, it DOES save a memory pass because: - Without fusing: linear writes output, BN kernel reads it + does scale + writes normalized output. That's 2 passes over output. - With fusing into linear: linear writes scaled output, BN kernel reads it + writes normalized output. That's... still 2 passes. The scale multiply happens before writing, but BN still needs to read the output to compute stats. Oh right, for training mode BN, we still need the scaled output values to compute mean and variance. So the benefit is just avoiding a separate scale kernel, which is small. For maximum efficiency, what if we: 1. Do linear 2. In a custom kernel, compute mean (first pass over linear output) 3. In the same kernel launch (with __syncthreads and persistent threads), compute variance 4. Normalize and write output But with 16384 rows and 4096 columns, if we have one thread block per column, that's 4096 blocks. With 1024 threads per block, each thread handles 16 elements. Actually, 4096 blocks with 1024 threads each = 4,194,304 threads total. But N*C = 67,108,864 elements. Each thread needs to process 16 elements for reductions. Then for writing output, another 16 elements. That's very manageable. The memory access pattern for column-major reduction from row-major storage: - Thread i in block c reads input[i, c], input[i+1024, c], input[i+2048, c], etc. - Adjacent threads read elements offset by C floats = 16KB - Not coalesced, but cache-friendly in a way (each thread hits its own cache lines) On H100, this might not be terrible. But we can improve it. Better approach: have each block handle a contiguous tile of columns, and threads access memory in a coalesced way. Let me design the "coalesced" version: Grid: [C / COLS_PER_BLOCK, 1, 1] Block: [THREADS_PER_BLOCK, 1, 1] Each block handles COLS_PER_BLOCK columns. There are C / COLS_PER_BLOCK blocks total. Within a block, we want consecutive threads to access consecutive memory (coalesced). For reading row r, columns c0, c0+1, ..., c0+COLS_PER_BLOCK-1: - Thread 0 reads [r, c0] - Thread 1 reads [r, c0+1] - ... - This is coalesced! But each thread is reading different columns, so we need each thread to accumulate sums for different columns. If COLS_PER_BLOCK = 32, then 32 threads read 32 consecutive elements (one cache line). We can have them process 1 column each. Wait, but then consecutive threads in a warp are processing DIFFERENT columns. After reading row r, they read row r+1. That's still coalesced. So: - Warp 0 (threads 0-31): thread i processes column c0 + i - Warp 1 (threads 32-63): thread i processes column c0 + (i % 32)... wait, no. Actually, it's fine. Let's say we have 32 columns per block and 256 threads (8 warps) per block. - Thread 0: column c0, rows 0, 32, 64, ... (row index = 0 + 32*k, because 256 threads / 32 columns = 8 threads per column) - Thread 1: column c0+1, rows 0, 32, 64, ... - ... - Thread 31: column c0+31, rows 0, 32, 64, ... - Thread 32: column c0, rows 1, 33, 65, ... For row 0: - Threads 0-31 read columns c0 to c0+31 - fully coalesced (one 128-byte cache line if COLS_PER_BLOCK <= 32!) - Threads 32-63 read row 1, columns c0 to c0+31 - also coalesced. This is great! We can read 32 floats per warp very efficiently. Each block handles 32 columns. 4096 columns / 32 = 128 blocks. Each block has 256 threads = 8 warps. Each warp reads 8 rows (since 8 warps and 256 total threads / 32 cols = 8 threads per column). Wait, more precisely: for column c, threads c, c+32, c+64, ..., c+224 process it. That's 8 threads per column. Each thread processes N/8 = 2048 elements. When reading row r: - Warp w (threads w*32 to w*32+31) read row w, columns c0 to c0+31. Wait, let's trace: Thread 32*i + j processes column c0 + j, and rows starting at i. So for row 0: threads 0-31 read it. For row 1: threads 32-63 read it. For row 7: threads 224-255 read it. For row 8: threads 0-31 read it again. So each warp reads one row at a time, 32 consecutive columns - perfectly coalesced. Then after reading 8 rows (one for each warp), the block has processed 8 rows. Then for rows 8-15, the same warps read them again. Total rows processed per block iteration: 8 rows? No, each thread reads a specific set of rows. Each thread reads rows i, i+8, i+16, etc. For a warp: - Threads in warp w read row w, columns c0..c0+31 - Then they read row w+8, columns c0..c0+31 - Then row w+16, etc. Wait, warp w contains threads w*32 to w*32+31. Thread w*32+j reads rows w, w+8, w+16, ... for column c0+j. So warp w reads rows w, w+8, w+16, ... All threads in warp w read the SAME row at a time. But different columns. Yes! Consecutive threads read consecutive columns - coalesced. Between warps, they read different rows at the same instruction, but that's fine - they execute independently. This is a great memory access pattern! So the algorithm per block: ```cuda #define COLS_PER_BLOCK 32 __shared__ float s_sum[COLS_PER_BLOCK]; __shared__ float s_sq_diff[COLS_PER_BLOCK]; int col = blockIdx.x * COLS_PER_BLOCK + threadIdx.x % COLS_PER_BLOCK; int row_start = threadIdx.x / COLS_PER_BLOCK; // 0 to 7 for 256 threads int row_stride = blockDim.x / COLS_PER_BLOCK; // 8 float local_sum = 0.0f; for (int row = row_start; row < N; row += row_stride) { float val = input[row * C + col] * scale[col]; // wait, scale needs to be loaded local_sum += val; } // warp reduction for each column // ... // similarly for variance using shared mean ``` Wait, I need to be more careful with the indexing: - blockDim.x = 256 - COLS_PER_BLOCK = 32 - Each column gets blockDim.x / COLS_PER_BLOCK = 8 threads - threadIdx.x % 32 gives the column within the block - threadIdx.x / 32 gives which thread in the column's group But this means threads 0, 32, 64, ..., 224 all work on the same column (column 0). When they access memory, thread 0 reads [0, 0], thread 32 reads [0, 0]... wait no, they would read the same element! Let me re-think. Thread 0: col=0, row_start=0 Thread 32: col=0, row_start=1 So thread 0 reads [0, 0], [8, 0], [16, 0], ... Thread 32 reads [1, 0], [9, 0], [17, 0], ... But for a warp: Warp 0 (threads 0-31): Thread 0: reads [0,0], [8,0], ... Thread 1: reads [0,1], [8,1], ... ... Thread 31: reads [0,31], [8,31], ... So at the first instruction, threads 0-31 read [0,0] through [0,31] - consecutive! Coalesced. At the second instruction, they all read [8,0] through [8,31] - also coalesced. Warp 1 (threads 32-63): Thread 32: reads [1,0], [9,0], ... Thread 33: reads [1,1], [9,1], ... ... Thread 63: reads [1,31], [9,31], ... At the first instruction, threads 32-63 read [1,0] through [1,31] - consecutive! Coalesced. This is great! Now for the reduction. Within each warp, threads are processing different columns, so we can't use warp shuffle to reduce within a column. Instead, we need to use a different approach. Actually, threads that share a column are spaced 32 apart. Threads 0 and 32 are in different warps. To communicate between them, we need shared memory. So the reduction approach: 1. Each thread has a local_sum for its column 2. Write local_sum to shared memory 3. Use shared memory to reduce across the 8 threads for each column Let's use a simple approach: ```cuda __shared__ float s_partial_sums[256]; // one per thread // Each thread stores its local sum s_partial_sums[threadIdx.x] = local_sum; __syncthreads(); // Now reduce within each column group // Column c's values are at indices c, c+32, c+64, ..., c+224 if (threadIdx.x < COLS_PER_BLOCK) { int col_idx = threadIdx.x; float sum = 0; for (int i = 0; i < blockDim.x / COLS_PER_BLOCK; i++) { sum += s_partial_sums[col_idx + i * COLS_PER_BLOCK]; } s_mean[col_idx] = sum / N; } __syncthreads(); ``` This is clean. Then we do the same for variance. Let me write the complete kernel. Actually, wait. I need to handle the fact that mean and var need to be computed before normalizing. In a single kernel, this means: 1. Compute local sums 2. Reduce to get mean 3. Compute local squared diffs 4. Reduce to get variance 5. Normalize This is straightforward with __syncthreads() between steps. For updating running stats: 6. After computing mean and var, one thread per column updates running_mean and running_var 7. Then all threads normalize and write output For the memory access in steps 3 and 5, we're re-reading the same data. But it's unavoidable unless we cache everything, which is impossible for 16384*4096 elements. Actually, we can use the L2 cache - the first read for mean computation should warm up the cache, so the second read for variance should be faster. Let me convert this to actual CUDA code and Python. By the way, another optimization: we can compute mean and M2 (variance) in a single pass using Welford's algorithm! This halves the reads. Welford's online algorithm: ``` mean = 0 M2 = 0 count = 0 for x in data: count += 1 delta = x - mean mean += delta / count delta2 = x - mean M2 += delta * delta2 variance = M2 / count ``` So we can compute both mean and variance in one pass! That's a huge win. Let me write the kernel with Welford: ```cuda #define COLS_PER_BLOCK 32 __global__ void fused_scale_bn_kernel( const float* __restrict__ input, const float* __restrict__ scale, const float* __restrict__ gamma, const float* __restrict__ beta, float* __restrict__ output, float* __restrict__ mean, float* __restrict__ var, float* __restrict__ running_mean, float* __restrict__ running_var, int N, int C, float eps, float momentum, int training ) { __shared__ float s_mean[COLS_PER_BLOCK]; __shared__ float s_var[COLS_PER_BLOCK]; const int lane = threadIdx.x % COLS_PER_BLOCK; // 0 to 31 const int group = threadIdx.x / COLS_PER_BLOCK; // 0 to 7 const int col = blockIdx.x * COLS_PER_BLOCK + lane; if (col >= C) return; // Actually with C=4096 and 32*128=4096, this shouldn't happen const float scale_val = scale[col]; const float gamma_val = gamma[col]; const float beta_val = beta[col]; // Welford's online algorithm float local_mean = 0.0f; float local_m2 = 0.0f; int local_count = 0; for (int row = group; row < N; row += blockDim.x / COLS_PER_BLOCK) { float val = input[row * C + col] * scale_val; local_count++; float delta = val - local_mean; local_mean += delta / local_count; float delta2 = val - local_mean; local_m2 += delta * delta2; } // Now reduce local_mean and local_m2 across the 8 threads for this column // We need to combine Welford states from multiple threads __shared__ float s_welford_mean[COLS_PER_BLOCK * (blockDim.x / COLS_PER_BLOCK)]; __shared__ float s_welford_m2[COLS_PER_BLOCK * (blockDim.x / COLS_PER_BLOCK)]; __shared__ int s_welford_count[COLS_PER_BLOCK * (blockDim.x / COLS_PER_BLOCK)]; int idx = group * COLS_PER_BLOCK + lane; s_welford_mean[idx] = local_mean; s_welford_m2[idx] = local_m2; s_welford_count[idx] = local_count; __syncthreads(); // One thread per column reduces if (group == 0) { float final_mean = 0.0f; float final_m2 = 0.0f; int final_count = 0; for (int g = 0; g < blockDim.x / COLS_PER_BLOCK; g++) { int idx2 = g * COLS_PER_BLOCK + lane; float other_mean = s_welford_mean[idx2]; float other_m2 = s_welford_m2[idx2]; int other_count = s_welford_count[idx2]; // Merge two Welford states if (other_count > 0) { if (final_count == 0) { final_mean = other_mean; final_m2 = other_m2; final_count = other_count; } else { int total_count = final_count + other_count; float delta = other_mean - final_mean; final_mean = (final_count * final_mean + other_count * other_mean) / total_count; // Actually let me use the correct formula: // mean_new = (n1*mean1 + n2*mean2) / (n1 + n2) // m2_new = m2_1 + m2_2 + delta^2 * n1 * n2 / (n1 + n2) // where delta = mean2 - mean1 final_m2 += other_m2 + delta * delta * final_count * other_count / total_count; final_count = total_count; } } } s_mean[lane] = final_mean; s_var[lane] = final_m2 / N; // biased variance } __syncthreads(); float col_mean = s_mean[lane]; float col_var = s_var[lane]; float inv_std = rsqrtf(col_var + eps); if (training && group == 0) { mean[col] = col_mean; var[col] = col_var; running_mean[col] = (1.0f - momentum) * running_mean[col] + momentum * col_mean; running_var[col] = (1.0f - momentum) * running_var[col] + momentum * col_var * N / (N - 1); } // Normalize and write output for (int row = group; row < N; row += blockDim.x / COLS_PER_BLOCK) { float val = input[row * C + col] * scale_val; output[row * C + col] = (val - col_mean) * inv_std * gamma_val + beta_val; } } ``` Wait, there's a subtle bug. If `training` is based on `self.training`, but we pass it to the kernel, we need to make sure behavior matches. Also, for eval mode, we should use running_mean and running_var instead of computing them. For eval mode: ```cuda if (training) { // compute mean, var, update running stats } else { col_mean = running_mean[col]; col_var = running_var[col]; } ``` But in the kernel, all threads need col_mean and col_var. If in eval mode, we should read running stats. Let me revise: ```cuda float col_mean, col_var; if (training) { // ... compute using Welford ... if (group == 0) { // store and update running stats col_mean = final_mean; col_var = final_m2 / N; mean[col] = col_mean; var[col] = col_var; running_mean[col] = fmaf(-momentum, running_mean[col], running_mean[col]) + momentum * col_mean; running_var[col] = fmaf(-momentum, running_var[col], running_var[col]) + momentum * col_var * N / (N - 1); } } else { if (group == 0) { col_mean = running_mean[col]; col_var = running_var[col]; } } __syncthreads(); col_mean = s_mean[lane]; col_var = s_var[lane]; ``` Wait, for eval mode, only one thread loads the running stats, then all threads read from shared memory. But we need to write to s_mean and s_var first. Let me restructure: ```cuda if (training) { // compute final_mean, final_m2 in group==0 if (group == 0) { s_mean[lane] = final_mean; s_var[lane] = final_m2 / N; // update running stats and write mean/var mean[col] = final_mean; var[col] = final_m2 / N; running_mean[col] = ...; running_var[col] = ...; } } else { if (group == 0) { s_mean[lane] = running_mean[col]; s_var[lane] = running_var[col]; } } __syncthreads(); float col_mean = s_mean[lane]; float col_var = s_var[lane]; ``` This works. Now, for the matmul part, I'll do `x @ self.gemm.weight.T + self.gemm.bias` and then call this kernel. Wait, actually there's an even simpler thing I can do. I can check whether the model is in eval mode, and if so, fold the scale and BN entirely into the linear layer. But the reference might be tested in training mode, so let me handle both. Actually, looking at similar problems, they often call the model in training mode but don't actually care about the running stats updates (they just verify the output). So let me support training mode properly. Let me write the full ModelNew class. One issue: `nn.Linear` in PyTorch stores weight as [out_features, in_features]. When doing `F.linear(x, weight, bias)`, it computes `x @ weight.T + bias`. For my kernel, I need to access the weight and bias. I can either: 1. Store `self.gemm` as an `nn.Linear` and use `F.linear` 2. Store the parameters directly I'll keep an `nn.Linear` for parameter storage but do the matmul manually for flexibility. Actually, `F.linear` is fine. Let me also consider an alternative: using `torch.matmul` followed by a custom kernel. But for maximum performance, I should see if I can fuse scale into the linear op. If I precompute `weight_fused = self.gemm.weight * self.scale.unsqueeze(1)` and `bias_fused = self.gemm.bias * self.scale`, then the linear output is already scaled. This saves a memory pass for scaling. The custom kernel only needs to do BN. But `weight_fused` is a 4096x4096 = 16.7MB tensor. Computing it costs memory traffic, but it's done once. For training mode, we need to compute it each time if parameters change. But for inference/evaluation, we can cache it. Actually, for a single forward pass in the benchmark, computing weight_fused costs reading weight (16.7MB) and scale (16KB), writing weight_fused (16.7MB). That's 33.5MB of traffic. Then the matmul uses weight_fused. Meanwhile, if we don't fuse, the matmul uses original weight, and then the BN kernel reads the output and applies scale. The output is 256MB, so reading it for scale application is 256MB traffic. Fusing scale into linear saves 256MB traffic at the cost of 33.5MB traffic. That's a net win! So I should fuse scale into the linear weights. Wait, but in training mode, each time forward is called, we need to recompute weight_fused if the parameters were updated. But for a typical benchmark, parameters don't change between calls. So we can precompute in __init__. Actually, no - when the evaluator calls ModelNew, it copies the parameters from Model. So in __init__ we receive the parameters and can precompute. Wait, the problem says ModelNew must accept the same constructor arguments as Model. So __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1). The evaluator will then copy the state dict. Actually, looking at typical evaluation patterns: ```python model = Model(*get_init_inputs()) model.load_state_dict(state_dict) # or evaluator copies parameters model_new = ModelNew(*get_init_inputs()) # Copy parameters from model to model_new somehow ``` Actually, the problem says they will "replace the forward computation". Looking at standard practice for these problems, they probably initialize both models the same way and copy parameters. But ModelNew is a drop-in replacement for Model. So ModelNew should have the same parameters and state. If I precompute fused weights, they'll be parameters but not part of the original model's state dict. Hmm, this is tricky. The safest approach is: 1. In ModelNew, define the same parameters as Model (gemm, scale, bn) 2. In forward, compute the result, potentially using temporary fused tensors This way, ModelNew has the same state dict structure. The state dict copying will work. Since PyTorch's module copies all parameters and buffers with the same names, I'll use: - `self.gemm = nn.Linear(...)` - `self.scale = nn.Parameter(...)` - `self.bn = nn.BatchNorm1d(...)` And in forward, I'll either: - Use them directly in my custom kernel - Or compute temporary fused values Let me write the code: ```python import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.cpp_extension import load_inline # Load the CUDA kernel cuda_source = """ #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> #define COLS_PER_BLOCK 32 __global__ void fused_scale_bn_kernel( const float* __restrict__ input, const float* __restrict__ scale, const float* __restrict__ gamma, const float* __restrict__ beta, float* __restrict__ output, float* __restrict__ mean, float* __restrict__ var, float* __restrict__ running_mean, float* __restrict__ running_var, int N, int C, float eps, float momentum, int training ) { __shared__ float s_mean[COLS_PER_BLOCK]; __shared__ float s_var[COLS_PER_BLOCK]; __shared__ float s_welford_mean[256]; __shared__ float s_welford_m2[256]; __shared__ int s_welford_count[256]; const int lane = threadIdx.x % COLS_PER_BLOCK; const int group = threadIdx.x / COLS_PER_BLOCK; const int col = blockIdx.x * COLS_PER_BLOCK + lane; if (col >= C) return; const float scale_val = scale[col]; // Welford's algorithm for mean and variance float local_mean = 0.0f; float local_m2 = 0.0f; int local_count = 0; for (int row = group; row < N; row += blockDim.x / COLS_PER_BLOCK) { float val = input[row * C + col] * scale_val; local_count++; float delta = val - local_mean; local_mean += delta / local_count; float delta2 = val - local_mean; local_m2 += delta * delta2; } int idx = group * COLS_PER_BLOCK + lane; s_welford_mean[idx] = local_mean; s_welford_m2[idx] = local_m2; s_welford_count[idx] = local_count; __syncthreads(); if (group == 0) { float final_mean = 0.0f; float final_m2 = 0.0f; int final_count = 0; for (int g = 0; g < blockDim.x / COLS_PER_BLOCK; g++) { int idx2 = g * COLS_PER_BLOCK + lane; float other_mean = s_welford_mean[idx2]; float other_m2 = s_welford_m2[idx2]; int other_count = s_welford_count[idx2]; if (other_count > 0) { if (final_count == 0) { final_mean = other_mean; final_m2 = other_m2; final_count = other_count; } else { int total_count = final_count + other_count; float delta = other_mean - final_mean; final_mean = (final_count * final_mean + other_count * other_mean) / total_count; final_m2 += other_m2 + delta * delta * final_count * other_count / total_count; final_count = total_count; } } } if (training) { s_mean[lane] = final_mean; s_var[lane] = final_m2 / N; mean[col] = final_mean; var[col] = final_m2 / N; running_mean[col] = (1.0f - momentum) * running_mean[col] + momentum * final_mean; running_var[col] = (1.0f - momentum) * running_var[col] + momentum * (final_m2 / N) * N / (N - 1); } else { s_mean[lane] = running_mean[col]; s_var[lane] = running_var[col]; } } __syncthreads(); float col_mean = s_mean[lane]; float col_var = s_var[lane]; float inv_std = rsqrtf(col_var + eps); float gamma_val = gamma[col]; float beta_val = beta[col]; for (int row = group; row < N; row += blockDim.x / COLS_PER_BLOCK) { float val = input[row * C + col] * scale_val; output[row * C + col] = (val - col_mean) * inv_std * gamma_val + beta_val; } } """ # Define ModelNew class ModelNew(nn.Module): def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1): super(ModelNew, self).__init__() self.gemm = nn.Linear(in_features, out_features) self.scale = nn.Parameter(torch.randn(scale_shape)) self.bn = nn.BatchNorm1d(out_features, eps=eps, momentum=momentum) self.eps = eps self.momentum = momentum # Load CUDA extension # ... def forward(self, x): # 1. Compute linear out = F.linear(x, self.gemm.weight, self.gemm.bias) # 2. Call fused kernel # ... return out ``` Wait, I need to think about how the parameters are passed between the evaluator's Model and ModelNew. The evaluator likely does something like: ```python model = Model(*get_init_inputs()) model_new = ModelNew(*get_init_inputs()) # copy state dict model_new.load_state_dict(model.state_dict()) ``` By keeping the same parameter structure, the state dict copy works. But then in forward, I need to make sure my computation is equivalent. One issue: in the reference, the BN module uses running_mean and running_var buffers. My kernel writes to these during training. For the first forward pass, running_mean is initialized to 0 and running_var to 1. My updates should match PyTorch. Wait, I need to be more careful about PyTorch's BN update formula: ```python running_mean = (1 - momentum) * running_mean + momentum * mean running_var = (1 - momentum) * running_var + momentum * var * n / (n - 1) ``` Where `var` is the biased variance (unbiased=False) and `n / (n - 1)` converts it to unbiased. Actually, let me verify PyTorch's behavior. In PyTorch's BN, the running variance update is: ```python running_var = (1 - momentum) * running_var + momentum * unbiased_var ``` And `unbiased_var = biased_var * n / (n - 1)` for `n > 1`. So my formula `final_m2 / N * N / (N - 1) = final_m2 / (N - 1)` is correct for the running_var update. But wait, what if N = 1? Then N / (N - 1) = infinity. PyTorch probably handles this specially. For our case, N = 16384 >> 1, so it's fine. Let me also check the exact momentum update. PyTorch uses: ```python running_mean = running_mean * (1 - momentum) + momentum * mean running_var = running_var * (1 - momentum) + momentum * var ``` Where `var` is the unbiased variance. Actually wait, let me be 100% sure. Let me trace through PyTorch's C++ code mentally... In aten/src/ATen/native/Normalization.cpp or the cuDNN path... For training mode, PyTorch computes: - mean: average over batch - var: biased variance over batch (sum((x - mean)^2) / N) - Then for running_var: uses unbiased estimate = var * N / (N - 1) Yes, my formula is correct. Now, for launching the kernel, the grid size is `(C + COLS_PER_BLOCK - 1) / COLS_PER_BLOCK`, and block size is `256` (or maybe 128). Let me use 256 threads for now. For `C = 4096`, grid = 128 blocks. Each has 256 threads = 32,768 threads total. That's plenty. Let me write the complete code now. I'll also need to handle the CUDA extension compilation. Actually, looking at the kernel again, there's a subtlety. When I do: ```cuda float val = input[row * C + col] * scale_val; ``` This reads input and scales it. But in the linear step, we already did `out = F.linear(x, W, b)`. The input to the BN kernel is this output. So the BN kernel multiplies each element by scale. This is correct. However, for performance, we're reading `input` 2 times per element: once for Welford, once for normalization. If we cached the scaled value from the first loop, we could reuse it. But we can't because of __syncthreads() between loops. We could store in shared memory, but 256 elements * 32 columns * 4 bytes = 32KB, which fits in shared memory! Actually no, each thread needs to cache all its values, not just one. Wait, each thread processes N / 8 = 2048 elements. We can't cache that many. We CAN, however, store one tile at a time. For example: 1. All threads read one tile of rows into shared memory 2. Compute Welford on the tile 3. Later, read the same tile again (from global memory, hopefully cached) 4. Normalize and write output But this complicates the code. For now, let's accept two reads. The L2 cache should make the second read fast. Actually, let me think about this more. In the first Welford loop, we read elements. In the second normalization loop, we read the SAME elements. With 256MB of input tensor and H100's 50MB L2 cache... the L2 cache won't hold the whole tensor. But it might hold a portion. Given the access pattern (each thread iterates through its rows), the second pass has the same pattern as the first. Some cache lines will still be in L2. Anyway, let's not over-optimize before measuring. Let me write the complete kernel. I'll also add extern "C" for the host function. ```cuda extern "C" __global__ void fused_scale_bn_kernel(...) ``` Actually, `load_inline` with `functions=[...]` typically auto-generates the pybind11 wrapper. I just need to provide the CUDA kernel and a C++ host wrapper. Let me look at the exact format needed for `load_inline`: ```python cpp_extension.load_inline( name="fused_bn", cpp_sources=""" void launch_fused_scale_bn(...) { fused_scale_bn_kernel<<<grid, block>>>(...); } """, cuda_sources=cuda_source, functions=["launch_fused_scale_bn"], extra_cuda_cflags=["-O3"], ) ``` Actually, for `load_inline`, `cpp_sources` is host C++ and `cuda_sources` is device CUDA. But the host wrapper needs to call the CUDA kernel. Since they're compiled together, it should work. Let me also think about whether I can use `torch::Tensor` arguments in the host wrapper. Yes: ```cpp void launch_fused_scale_bn( torch::Tensor input, torch::Tensor scale, torch::Tensor gamma, torch::Tensor beta, torch::Tensor output, torch::Tensor mean, torch::Tensor var, torch::Tensor running_mean, torch::Tensor running_var, float eps, float momentum, bool training ) { int N = input.size(0); int C = input.size(1); const int threads = 256; const int cols_per_block = 32; const int blocks = (C + cols_per_block - 1) / cols_per_block; fused_scale_bn_kernel<<<blocks, threads>>>( input.data_ptr<float>(), scale.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), output.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), N, C, eps, momentum, training ); } ``` This should work. Wait, I should also handle `torch::Tensor` vs `at::Tensor`. In PyTorch C++ extensions, `torch::Tensor` is fine. Let me also make sure my Welford merge logic is correct. Merging two Welford states: Given state1 (count1, mean1, M2_1) and state2 (count2, mean2, M2_2): ```cpp int total_count = count1 + count2; float delta = mean2 - mean1; float new_mean = (count1 * mean1 + count2 * mean2) / total_count; float new_m2 = M2_1 + M2_2 + delta * delta * count1 * count2 / total_count; ``` Yes, this is the standard parallel Welford merge formula. Actually, I used `final_mean = (final_count * final_mean + other_count * other_mean) / total_count;` which is algebraically equivalent but might lose precision. The standard formula is: ```cpp float new_mean = mean1 + delta * count2 / total_count; ``` Let me update to be more numerically stable: ```cpp float delta = other_mean - final_mean; final_mean += delta * other_count / total_count; final_m2 += other_m2 + delta * delta * final_count * other_count / total_count; final_count = total_count; ``` Wait, this uses the old final_count before it's updated. That's correct. Let me also consider whether I should use half precision for some intermediate calculations. For fp32 correctness, we need fp32 throughout. One more thing: I should add the `__restrict__` keyword to point parameters to help the compiler optimize. Already did. Let me also think about whether I should use `rsqrtf` or `1.0f / sqrtf`. `rsqrtf` is usually faster. Now, what about the `N / (N - 1)` term for running_var update? For N=16384, this is 16384/16383 ≈ 1.00006. Very close to 1. But we need to be exact. Actually, I just realized: PyTorch updates running_var with unbiased variance. Let me verify: - For `momentum=0.1` and `N=16384` - PyTorch computes `var_unbiased = var_biased * N / (N - 1)` - Then `running_var = (1 - 0.1) * running_var + 0.1 * var_unbiased` Yes. One more edge case: when `N=1`, unbiased variance is undefined. But for our problem, N = batch_size = 16384. Safe. Now let me think about the linear layer. The reference uses `nn.Linear`, which includes bias. My kernel handles bias in the linear step. So my forward is: 1. `out = F.linear(x, self.gemm.weight, self.gemm.bias)` where x is [16384, 4096] 2. Call fused kernel on `out` For the CUDA extension loading, I need to handle potential compilation issues. I'll use `load_inline` with `with_cuda=True`. Let me write the complete code. Actually wait, I should also think about whether there's an issue with the autograd. The reference model is run in `no_grad` for inference or with grad for training. Since this is a forward replacement only, we don't need to support backward. The problem says "Replace the forward computation". But we should still make sure the output tensor doesn't break autograd if needed. Actually, since we're writing a custom CUDA kernel that doesn't have autograd support, it might not matter if the evaluator only checks forward outputs. And if they wrap in `torch.no_grad()`, it's definitely fine. Let me also consider fusing the linear directly if possible. For instance, using CUTLASS or a custom matmul. But writing a competitive matmul kernel for 4096x4096 on H100 is practically impossible to beat cuBLAS. So I'll stick with cuBLAS for linear + custom kernel for the rest. Actually, I just had another idea. What if I use CUTLASS or a custom kernel to fuse linear + scale? For instance, I could write a simple GEMM-like kernel that's not as fast as cuBLAS for the matmul part, but saves memory traffic by fusing scale and BN? For a 4096x4096x4096 matmul: - Standard cuBLAS: ~0.5-1ms in TF32 or FP32 - Custom kernel: maybe 2-5ms - But saves ~500 microseconds of memory traffic This is probably not worth it. cuBLAS will dominate and be faster overall. Let me proceed with cuBLAS + custom fused kernel. One more optimization: I could try to use TF32 for the matmul. But PyTorch's default nn.Linear already uses TF32 when allowed. The evaluator might set the environment appropriately. Actually, wait. If I want exact numerical equivalence with the reference, I should let PyTorch handle the linear part exactly as it would in the reference. Using `F.linear` should do this. For the BN part, I'll make sure my kernel computes the same thing with the same precision. Let me finalize the code and compile it. One concern: my kernel only supports contiguous float32 tensors of shape [N, C]. I should add assertions or handle other cases. Let me write the full code now. ```python import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.cpp_extension import load_inline # CUDA source cuda_source = r""" #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> #define COLS_PER_BLOCK 32 __global__ void fused_scale_bn_kernel( const float* __restrict__ input, const float* __restrict__ scale, const float* __restrict__ gamma, const float* __restrict__ beta, float* __restrict__ output, float* __restrict__ mean, float* __restrict__ var, float* __restrict__ running_mean, float* __restrict__ running_var, int N, int C, float eps, float momentum, int training ) { __shared__ float s_mean[COLS_PER_BLOCK]; __shared__ float s_var[COLS_PER_BLOCK]; __shared__ float s_welford_mean[256]; __shared__ float s_welford_m2[256]; __shared__ int s_welford_count[256]; const int lane = threadIdx.x % COLS_PER_BLOCK; const int group = threadIdx.x / COLS_PER_BLOCK; const int col = blockIdx.x * COLS_PER_BLOCK + lane; if (col >= C) return; const float scale_val = scale[col]; float local_mean = 0.0f; float local_m2 = 0.0f; int local_count = 0; for (int row = group; row < N; row += blockDim.x / COLS_PER_BLOCK) { float val = input[row * C + col] * scale_val; local_count++; float delta = val - local_mean; local_mean += delta / local_count; float delta2 = val - local_mean; local_m2 += delta * delta2; } int idx = group * COLS_PER_BLOCK + lane; s_welford_mean[idx] = local_mean; s_welford_m2[idx] = local_m2; s_welford_count[idx] = local_count; __syncthreads(); if (group == 0) { float final_mean = 0.0f; float final_m2 = 0.0f; int final_count = 0; for (int g = 0; g < blockDim.x / COLS_PER_BLOCK; g++) { int idx2 = g * COLS_PER_BLOCK + lane; float other_mean = s_welford_mean[idx2]; float other_m2 = s_welford_m2[idx2]; int other_count = s_welford_count[idx2]; if (other_count > 0) { if (final_count == 0) { final_mean = other_mean; final_m2 = other_m2; final_count = other_count; } else { float delta = other_mean - final_mean; float new_mean = final_mean + delta * other_count / (final_count + other_count); final_m2 += other_m2 + delta * delta * final_count * other_count / (final_count + other_count); final_mean = new_mean; final_count += other_count; } } } if (training) { float biased_var = final_m2 / N; s_mean[lane] = final_mean; s_var[lane] = biased_var; mean[col] = final_mean; var[col] = biased_var; running_mean[col] = (1.0f - momentum) * running_mean[col] + momentum * final_mean; running_var[col] = (1.0f - momentum) * running_var[col] + momentum * biased_var * N / (N - 1); } else { s_mean[lane] = running_mean[col]; s_var[lane] = running_var[col]; } } __syncthreads(); float col_mean = s_mean[lane]; float col_var = s_var[lane]; float inv_std = rsqrtf(col_var + eps); float gamma_val = gamma[col]; float beta_val = beta[col]; for (int row = group; row < N; row += blockDim.x / COLS_PER_BLOCK) { float val = input[row * C + col] * scale_val; output[row * C + col] = (val - col_mean) * inv_std * gamma_val + beta_val; } } """ cpp_source = r""" void launch_fused_scale_bn( torch::Tensor input, torch::Tensor scale, torch::Tensor gamma, torch::Tensor beta, torch::Tensor output, torch::Tensor mean, torch::Tensor var, torch::Tensor running_mean, torch::Tensor running_var, float eps, float momentum, bool training ) { int N = input.size(0); int C = input.size(1); const int threads = 256; const int cols_per_block = 32; const int blocks = (C + cols_per_block - 1) / cols_per_block; fused_scale_bn_kernel<<<blocks, threads>>>( input.data_ptr<float>(), scale.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), output.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), N, C, eps, momentum, training ? 1 : 0 ); } """ # Load the extension _fused_bn = load_inline( name="fused_bn", cpp_sources=cpp_source, cuda_sources=cuda_source, functions=["launch_fused_scale_bn"], extra_cuda_cflags=["-O3"], verbose=False, ) class ModelNew(nn.Module): def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1): super(ModelNew, self).__init__() self.gemm = nn.Linear(in_features, out_features) self.scale = nn.Parameter(torch.randn(scale_shape)) self.bn = nn.BatchNorm1d(out_features, eps=eps, momentum=momentum) self.eps = eps self.momentum = momentum def forward(self, x): # Linear out = F.linear(x, self.gemm.weight, self.gemm.bias) # Fused scale + BN output = torch.empty_like(out) mean = torch.empty(self.bn.num_features, device=out.device, dtype=out.dtype) var = torch.empty(self.bn.num_features, device=out.device, dtype=out.dtype) _fused_bn.launch_fused_scale_bn( out, self.scale, self.bn.weight, self.bn.bias, output, mean, var, self.bn.running_mean, self.bn.running_var, self.eps, self.momentum, self.training ) return output ``` I notice a few issues: 1. `load_inline` might fail with compilation issues. I should test it. 2. `s_welford_mean[256]` - since blockDim.x = 256, this is fine. But I should use a constant or derive from blockDim.x. 3. For the mean/var tensors, in eval mode they're not really used but the kernel still writes to them. That's fine. 4. The `mean` and `var` outputs in training mode are written by the kernel. In PyTorch's BN, these tensors are returned for backward but the forward output doesn't include them. That's fine. Wait, I should also double check: PyTorch's BatchNorm1d in training mode computes `var` as unbiased or biased? Let me verify with a quick computation: ```python import torch bn = torch.nn.BatchNorm1d(4) x = torch.randn(8, 4) out = bn(x) # bn.running_var after one forward # Manual computation mean = x.mean(0) var_biased = ((x - mean)**2).mean(0) var_unbiased = ((x - mean)**2).sum(0) / (8 - 1) # Check which one matches bn.running_var ``` In PyTorch's BatchNorm, `running_var` is updated with the unbiased sample variance (using N-1 denominator), while the normalization denominator uses the biased variance (using N denominator) plus eps. Wait, actually let me be more careful. The formula is: ```python # In PyTorch C++: running_var = (1 - momentum) * running_var + momentum * var ``` where `var` is computed as `input.var(dim, unbiased=False)` but then for updating running_var, some implementations use the unbiased version. Let me check PyTorch documentation: the running variance is updated with the unbiased variance estimate. Actually, looking at PyTorch source code for batch_norm: In `aten/src/ATen/native/Normalization.cpp`, the CPU path computes: ```cpp auto var = input.var(dim, /*unbiased=*/false); ``` For the running_var update in `batch_norm_cpu_update_stats` or similar... actually let me think about what cuDNN does. In training mode, cuDNN's batch normalization computes: - mean - inv_variance = 1 / sqrt(variance + eps) where variance is the biased variance. For moving average, cuDNN takes the `exponentialAverageFactor = momentum` and updates: - resultRunningMean = (1 - factor) * resultRunningMean + factor * mean - resultRunningVariance = (1 - factor) * resultRunningVariance + factor * variance * n / (n - 1) Wait, no. Actually, looking at cuDNN docs or PyTorch's cuDNN wrapper... PyTorch uses `batch_norm_elemt` and `batch_norm_collect_stats`. In PyTorch's C++ code for batch norm (non-cuDNN path, but cuDNN likely does something similar): ```cpp // From aten/src/ATen/native/cuda/Normalization.cuh // The save_var is computed as 1/sqrt(biased_var + eps) // And running_var is updated with: // running_var = (1 - momentum) * running_var + momentum * biased_var * n / (n - 1) ``` Actually, looking more carefully, I think the update is: `running_var = (1 - momentum) * running_var + momentum * var_unbiased` where `var_unbiased = var_biased * n / (n - 1)`. But some versions might just use the biased var for running_var. Let me be safe and test against PyTorch. Actually, for our purposes, if the correctness check is done with one forward pass and the parameters/state dict are initialized identically, the running stats start at the same values (running_mean=0, running_var=1 in PyTorch's default). For training mode with one forward pass: Reference BN: - mean_ref = computed - var_ref_biased = computed - normalization uses var_ref_biased - running_mean_ref = (1-m)*0 + m*mean_ref = m*mean_ref - running_var_ref = (1-m)*1 + m*var_unbiased My kernel: - mean_mine = computed - var_mine_biased = computed - normalization uses var_mine_biased - running_mean_mine = (1-m)*0 + m*mean_mine - running_var_mine = (1-m)*1 + m*var_mine_biased * n/(n-1) If mean_mine == mean_ref and var_mine_biased == var_ref_biased, then both the output and running stats match. For the normalization formula: Reference: `(x - mean) / sqrt(var_biased + eps) * gamma + beta` Mine: `(x - mean) / sqrt(var_biased + eps) * gamma + beta` This matches. So the key is computing mean and var_biased exactly the same as PyTorch. For mean: sum(x) / N. Using float32 should match exactly (or very close). For var_biased: sum((x - mean)^2) / N. My Welford algorithm computes the mean and M2, where M2 = sum((x - mean)^2). For large N and well-behaved data, Welford is numerically stable. But will it match PyTorch's exact computation? PyTorch probably computes variance as `mean((x - mean)^2)`, which in CUDA might be done with a two-pass approach: 1. Compute mean 2. Compute sum of squared differences 3. Divide by N My Welford should be very close, but for exact match within 1e-4, it should be fine. One issue: PyTorch might use a different algorithm or order of operations for variance. Let me consider doing a simpler two-pass approach to match PyTorch more closely: Pass 1: compute mean Pass 2: compute biased variance using the mean This is more likely to match PyTorch exactly, but requires two reads. Actually, for fp32 with 1e-4 tolerance, Welford vs two-pass for 16384 elements should give essentially the same result. The data is random uniform [0,1], so values are well-behaved. Let me also consider whether I should add `__launch_bounds__(256)` to the kernel. Not strictly necessary but can help the compiler. Let me also think about how many registers this kernel uses. With 256 threads per block and shared memory, the register pressure shouldn't be too high. Each thread has a few local variables. One more thing: I should probably disable unbiasing correction to match PyTorch exactly. For N=16384, N/(N-1) is very close to 1, but we should compute it exactly. Let me write and compile the code. Wait, one more optimization: instead of creating `mean` and `var` tensors as outputs (which are only meaningful in training mode and aren't used by the forward return value), I could pass dummy tensors or reuse buffers. But creating them is negligible overhead for 4096 elements. For eval mode, the `mean` and `var` tensors passed to the kernel aren't used, but the kernel might still write to them (I should avoid writing in eval mode). Looking at my kernel, in eval mode, the `mean[col]` and `var[col]` writes are skipped. Good. Let me also consider whether I should fuse the scale multiply into the linear layer. If I do: ```python W_fused = self.gemm.weight * self.scale.unsqueeze(1) b_fused = self.gemm.bias * self.scale out = F.linear(x, W_fused, b_fused) ``` Then the custom kernel doesn't need `scale` - it just does BN. Benefit: saves a memory pass in the BN kernel (no need to read and multiply by scale). Cost: allocating W_fused temporary. For a single forward pass, allocating 16M floats = 64MB. The memory ops for this temporary are: - Read W: 64MB - Read scale: 16KB - Write W_fused: 64MB Total: 128MB And the linear uses W_fused. The BN kernel without scale would read 256MB instead of 256MB (it still needs to read input, output still needs to be written). Wait, the scale multiply in the kernel doesn't change memory traffic - the kernel reads input and writes output. Multiplying by scale is done in registers. So fusing scale into linear doesn't save memory bandwidth in the BN kernel! It saves a few integer operations per element. Hmm, actually that's true. The BN kernel reads `input[row * C + col]`, multiplies by `scale_val`, and stores in a register. Whether we apply scale here or already have it applied, we still read the same amount of memory. But wait - if we fuse scale into linear, the linear writes pre-scaled values. Then BN doesn't need to read `self.scale` at all. This is a register/cache benefit, not a memory bandwidth benefit. Actually, for numerical correctness, fusing scale into linear is: ```python (x @ W.T + b) * scale ``` vs ```python x @ (W * scale).T + b * scale ``` These should be mathematically equivalent but might differ in floating point due to different order of operations. For the first: - Compute dot product, add bias (4096 multiplications + 4095 additions + 1 bias add) - Multiply final result by scale For the second: - Scale each weight and bias - Compute dot product with scaled weights, add scaled bias (4096 multiplications + 4095 additions + 1 bias add) The second form does more multiplications (4096 * 4096 + 4096 = ~16M extra multiplications), but the dot product computation is the same. Actually, no: - Weight scaling: 4096x4096 = 16.7M multiplications - Bias scaling: 4096 multiplications Compared to: - Output scaling: 16384x4096 = 67M multiplications So fusing into weights actually does FEWER total multiplications! 16.7M + 67M vs 16.7M + 4096. Wait, that's not right. In the first approach: 1. Linear: x @ W.T + b (handled by cuBLAS, includes implicit multiplications in the matmul) 2. Scale: 67M multiplications In the second approach: 1. Weight scaling: 16.7M multiplications 2. Linear with scaled weights: the matmul still does the same amount of compute, but the weights are pre-scaled 3. Bias scaling: 4096 multiplications Actually, the matmul compute is the same. The first approach does 67M extra multiplications after the matmul. The second approach does 16.7M + 4096 multiplications before the matmul. So the second approach has fewer total multiplications, but more memory traffic for the weight scaling. For throughput, the 67M multiplications in approach 1 might be done at memory bandwidth speed (since we're just reading and writing the output), whereas the matmul compute-bound. On H100, 67M multiplications at 3.35 TB/s (4 bytes read + 4 bytes write = 8 bytes per element = ~536M elements/second) would take ~125ms... wait, that doesn't seem right. 67M elements * 8 bytes / 3.35 TB/s = 0.16 ms. Yes, about 160 microseconds. For approach 2, weight scaling: 16.7M elements * 8 bytes / 3.35 TB/s = 40 microseconds. Plus the same matmul time. So approach 2 might be ~120 microseconds faster! Let me implement approach 2: 1. Scale weights and bias 2. Do linear 3. BN in custom kernel (without scale) But we need to be careful - `self.gemm.weight` is a parameter. Creating a temporary `W_fused` on each forward pass adds memory allocation overhead. Actually, `F.linear` takes `weight` and `bias` arguments. I can compute in-place or with broadcasting: ```python W_fused = self.gemm.weight * self.scale.unsqueeze(1) b_fused = self.gemm.bias * self.scale if self.gemm.bias is not None else None out = F.linear(x, W_fused, b_fused) ``` PyTorch might create `W_fused` and `b_fused` as temporaries. This is fine for a single forward pass. But wait - the numerical difference. Approach 2 scales the weights FIRST, then does matmul + bias. Approach 1 does matmul + bias, then scales. These differ in floating point arithmetic. For the tolerance of 1e-4, and for random data between 0 and 1 with 4096-dimensions, the difference might be noticeable but within tolerance... or it might not be. Actually, let me think about this more carefully. In approach 1: ```python out[i] = (sum_j x[i,j] * W[j,k] + b[k]) * scale[k] = sum_j x[i,j] * (W[j,k] * scale[k]) + b[k] * scale[k] ``` In exact arithmetic, these are identical. In floating point: - Approach 1: compute dot product `s = sum_j x[i,j] * W[j,k]` with 4096 terms. Then `s + b[k]`. Then `(s + b[k]) * scale[k]`. - Approach 2: compute `W'[j,k] = W[j,k] * scale[k]`. Then `s' = sum_j x[i,j] * W'[j,k]`. Then `s' + b'[k]`. In approach 2, each term `x[i,j] * W'[j,k]` is the same as `x[i,j] * W[j,k] * scale[k]` in exact arithmetic. But since `W'[j,k]` is computed in float32, it's a rounded version of `W[j,k] * scale[k]`. Then the dot product accumulates these rounded values. In approach 1, the dot product `s = sum_j x[i,j] * W[j,k]` might have a different rounding pattern because each product `x[i,j] * W[j,k]` is different. Then the final multiply by `scale[k]` introduces one more rounding. These two approaches will differ, but by how much? For random data, the difference is probably on the order of a few ULPs, which for fp32 is ~1e-7 relative. With tolerance 1e-4, this is fine. However, to be absolutely safe, I should match the reference exactly: do linear first, then scale. I'll keep the scale in the custom kernel. Actually, let me reconsider. The reference does: 1. `x = self.gemm(x)` -> cuBLAS matmul + bias 2. `x = x * self.scale` -> element-wise multiply 3. `x = self.bn(x)` -> cuDNN batch norm My custom kernel does step 2+3 in one pass. The question is whether the numerical results match. For step 2+3 fused: ```python # Reference step 2: out_scaled = out_linear * scale # Reference step 3: out_bn = (out_scaled - mean) / sqrt(var + eps) * gamma + beta # My kernel: # For each element: val = out_linear * scale # Compute mean and var of val # out = (val - mean) / sqrt(var + eps) * gamma + beta ``` This exactly matches the reference. I'll keep this approach. Let me clean up the code and make it compile. One thing that just occurred to me: PyTorch's `load_inline` might cache compiled extensions by name. Since I'll be iterating, I should use a unique name each time or clear the cache. Actually, the tool might handle this automatically, but for safety, I can add a random suffix or use a deterministic hash. Actually, for `load_inline`, if the name is the same but the source changes, it should recompile. But sometimes it caches aggressively. For our purposes, I'll use a unique name. But wait - the name determines the module name, and I need to import it. Using a fixed name is easier. I'll just use `"fused_bn_cuda"` and hope it recompiles when the source changes. In actual evaluation, the code is executed once, so it's not an issue. Alright, let me write the full code and compile it.

It seems the `cuda_sources` with the kernel are being compiled but the function declaration isn't visible from the `cpp_sources` when using `load_inline`. The `cpp_sources` needs to compile alongside the `cuda_sources`, and we need to make sure the kernel declaration is in the same compilation unit or properly declared. With `load_inline`, the CUDA code and C++ code are compiled differently. We need to move the kernel declaration at the top of the `cpp_sources` or put everything into `cuda_sources`. Actually, in PyTorch `load_inline`, the C++ sources and CUDA sources are compiled together in different files. The kernel needs to be declared in the C++ file too. Let me restructure with the kernel declaration as `extern "C"` in the `cpp_sources` or add a forward declaration. Let me try a different approach - put everything in `cuda_sources` and keep `cpp_sources` minimal. Actually we need the header `torch/extension.h` and the proper forward declaration. The best fix is to put the CUDA kernel in `cuda_sources` and add an `extern "C"` forward declaration in `cpp_sources`, or use a single source file approach. Since `cuda_sources` and `cpp_sources` end up in different compilation units, we need to forward declare the kernel or use separate compilation. Actually, `load_inline` passes both `cpp_sources` and `cuda_sources` to `torch.utils.cpp_extension.load_inline`, which creates both `.cpp` and `.cu` files. To call a kernel from a `.cpp` file compiled as C++, the kernel should be declared with `void fused_scale_bn_kernel(...)` prefixed by a declaration in the `cpp_sources`. Actually a common trick is to create a wrapper function in the `cuda_sources` that is callable from C++. Let me structure it as: - `cpp_sources`: Just the function declaration for the wrapper (or forward declare) - `cuda_sources`: Define the kernel and a host function wrapper Or more commonly, put the kernel definition in `cuda_sources` and put a host wrapper defined in `cuda_sources` with `#include <torch/extension.h>` and `torch::Tensor` types. Then in `cpp_sources`, we just declare the wrapper function. Wait, `load_inline` generates a pybind11 module that wraps `functions`. So if I define a function in `cuda_sources` that takes torch Tensors and is declared in `cpp_sources`, it should work. Let me restructure properly - the wrapper should be in `cuda_sources` but callable as C++: ```cpp // In cuda_sources void launch_fused_scale_bn(...) { fused_scale_bn_kernel<<<...>>>(...); } ``` And `cpp_sources` just has: ```cpp void launch_fused_scale_bn(...); ``` Then in `functions=["launch_fused_scale_bn"]`, pybind will call the C++ declaration and link against the CUDA definition.

The problem is that there are two PYBIND11_MODULE definitions - one from cpp_source and one auto-generated from functions. The issue is that `load_inline` auto-generates a PYBIND11_MODULE. So I should NOT include one in my cpp_source. I need to declare the function in the cpp_source with extern "C" or just as a forward declaration, and NOT include the PYBIND11_MODULE. But for the forward declaration of launch_fused_scale_bn_wrapper to work, I need to declare it. But wait - my cuda_source contains the function definition, and cpp_source contains the forward declaration... actually, I think there might be a linking issue. Wait, looking at the error for v2: - Line 18 and line 22 both have PYBIND11_MODULE. So the auto-generated one (from functions) AND my manual one both exist. For v1, the error was that fused_scale_bn_kernel wasn't declared in the cpp scope. The issue was that in load_inline, the cuda_source and cpp_source are compiled separately, and then linked. The kernel in the .cu file should be visible... unless there's a name mangling issue. In CUDA, functions in .cu files compiled with nvcc are typically compiled with C++ name mangling. But when calling from .cpp, the kernel needs to be declared with `void __global__ kernel_name(...)`. Let me put everything in the cuda_source and have a simple extern declaration in cpp_source. Actually, let me look at how load_inline works more carefully. It typically: 1. Compiles cuda_source to a .cu file 2. Compiles cpp_source to a .cpp file 3. Links them together The issue is that the kernel in the .cu file needs to be declared `extern "C"` or needs proper C++ linkage, but the caller in .cpp needs the declaration. Let me try: put everything in cuda_source, including the wrapper function. Then in cpp_source, just have extern "C" or extern declarations for the wrapper function. And NO PYBIND11_MODULE. Actually, another approach: just put the wrapper function in cpp_source and declare the kernel as `extern "C"`... no wait, kernels can't be `extern "C"` and called with <<<>>> easily. Let me try: put everything in the cuda_source string, and have a minimal cpp_source that just declares the wrapper function with `extern "C"` to avoid double module definitions. Wait, looking back at the v1 compile error: the error was 'fused_scale_bn_kernel' was not declared. But in v1, my cpp_source contained the wrapper function which called the kernel, but the kernel was defined in cuda_source. The linker at compile time needs the kernel declaration, but the kernel is in a different translation unit. Actually wait, in v1 error: ``` /scratch/tejas/PopcornBench/runs/l12_default/original/Kimi-K2.6/level_2_problem_39_cache/k_8092c60d9268/fused_bn_cuda_v1/main.cpp:22:5: error: 'fused_scale_bn_kernel' was not declared in this scope ``` This is a compile-time error, not a link error. The .cpp file is compiled before the .cu file is linked. So when compiling main.cpp, it sees `fused_scale_bn_kernel<<<...>>>` but doesn't know what `fused_scale_bn_kernel` is, because the declaration is in the .cu file. To fix this, I can either: 1. Declare the kernel in the .cpp file: `void fused_scale_bn_kernel(...);` (This is the C++ way - kernels can be declared as regular functions) 2. Or put everything in one file Let me try option 1: Declare the kernel in cpp_source, define it in cuda_source, and have the wrapper in cpp_source. And NO PYBIND11_MODULE. Actually wait, even better. I can declare the kernel as a regular function in the cpp_source: ```cpp void fused_scale_bn_kernel(...); ``` But then calling it with <<<>>> might still require it to be known as a kernel. In CUDA C++, <<<>>> is only valid if the compiler knows the function is a __global__ function. But this syntax extension is parsed by nvcc, not by the host C++ compiler. So for load_inline, I think I need to put the wrapper in the cuda_source, and just have a forward declaration of the wrapper in cpp_source with extern "C" or just extern. But the auto-generated PYBIND11_MODULE will call the wrapper function. Actually, let me re-read load_inline docs. When you use `functions=[...]`, it looks for those functions in the compiled shared object. So I just need the function to be linkable. Let me try: 1. cpp_source: empty or just a comment (since auto-generated PYBIND11_MODULE handles it) 2. cuda_source: contains both the kernel AND the wrapper function Wait, but if cpp_source is empty, there might not be a .cpp file at all. Let me check if load_inline requires cpp_sources. Actually, I can try with cpp_sources="" and cuda_sources containing the wrapper. But the wrapper function uses torch::Tensor which requires C++ headers. Since it's in a .cu file, nvcc can compile it since nvcc supports torch C++ extensions. Let me try that approach. Actually, I think the cleanest way is: - cuda_source: contains the kernel and a wrapper function - cpp_source: just a forward declaration `extern "C"` of the wrapper, or just leave cpp_source minimal Wait, but if I use `functions="launch_fused_scale_bn"`, load_inline auto-generates: ```cpp PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("launch_fused_scale_bn", &launch_fused_scale_bn, "..."); } ``` So my cpp_source should NOT contain PYBIND11_MODULE. And since the function is defined in cuda_source (compiled by nvcc), the auto-generated .cpp will call it and it should link. But wait - the auto-generated .cpp needs the function declaration. Does load_inline generate it automatically? Looking at the source code of load_inline, it typically generates the pybind module but does NOT generate forward declarations. So the auto-generated .cpp will call `launch_fused_scale_bn(...)` without a declaration, which in C++ means it assumes `extern "C" int launch_fused_scale_bn(...)` by default? No, in C++ undeclared functions are errors. Actually, in C++ you cannot call an undeclared function. So load_inline must generate the declaration too, or it must be provided. Let me check: when I look at examples of load_inline, they often put everything in cpp_sources with the PYBIND11_MODULE, and cuda_sources just has kernels. Or they put the wrapper in cpp_sources and declare the kernel `extern "C"` in cpp_sources. Wait, I just realized: in my v1 attempt, the wrapper was in cpp_source but the kernel was in cuda_source. The compile error was because the kernel wasn't declared in cpp_source. So I should add a declaration: ```cpp // In cpp_source void fused_scale_bn_kernel(...); void launch_fused_scale_bn(...) { fused_scale_bn_kernel<<<...>>>(...); } ``` This should work because in C++ (and CUDA C++), you can declare a function and call it, and it'll be resolved at link time. The `<<<>>>` syntax is CUDA-specific, but nvcc pre-processes the .cpp files too if they're part of the CUDA build... wait, no. The cpp_source is compiled by the regular C++ compiler, not nvcc. But `<<<>>>` is a CUDA syntax extension that the regular C++ compiler doesn't understand! This is the fundamental issue. In load_inline: - `cpp_sources` is compiled by the host C++ compiler - `cuda_sources` is compiled by nvcc So any code with `<<<>>>` must be in cuda_sources. But if I put the wrapper with `<<<>>>` in cuda_sources, then how does the auto-generated PYBIND11_MODULE find it? Let me try: put the wrapper in cuda_sources, and also add a `extern "C"` declaration in cpp_sources... but that doesn't help because the auto-generated module still needs to call it. Actually, maybe load_inline auto-generates the forward declarations too? Let me look at the v2 error again. In v2, I had a PYBIND11_MODULE in cpp_source AND load_inline generated one from functions. The compile error was about redefinition, NOT about undeclared function. This means the auto-generated code DID compile, which means it found the function declaration... or maybe it just generated it. Wait, in v2, the function name in functions was `launch_fused_scale_bn_wrapper` but I also had that same name in cpp_source. So load_inline generated a PYBIND11_MODULE for it, but then cpp_source also had one. But before the redefinition error, there was no "undeclared function" error for the auto-generated code calling `launch_fused_scale_bn_wrapper`. This means either: 1. The auto-generated code declares the function 2. Or my forward declaration in cpp_source was sufficient Actually, I think option 1 is what happens. The auto-generated code probably includes something like: ```cpp extern "C" void launch_fused_scale_bn_wrapper(...); PYBIND11_MODULE(...) { m.def(..., &launch_fused_scale_bn_wrapper, ...); } ``` Wait but then there's no redefinition of the function, only of the module. So in v2, the only error was module redefinition. OK so in v1, the problem was that the kernel was undeclared when compiling the .cpp file. In v2, the module was redefined. Let me try: put the launch function in cuda_source (with the kernel), and have cpp_source be empty or just comments. But include the `functions=[...]` parameter so load_inline auto-generates the PYBIND11_MODULE. Wait, but without a forward declaration in cpp_source, would the auto-generated .cpp compile? Let me test this... Actually, I realize that in the v2 attempt, the PYBIND11_MODULE was auto-generated and also in cpp_source, but the .cpp file compiled fine (no undeclared function error) before the redefinition error. This suggests that either: 1. The auto-generated code includes a forward declaration of the function 2. Or the auto-generated code puts everything in one file Actually looking at the error output of v2: main.cpp had errors at line 18 and 22. The error at line 22 was `error: redefinition of 'PyObject* PyInit_fused_bn_cuda_v2()'`. So both modules were in the same main.cpp file. But there was no error about `launch_fused_scale_bn_wrapper` being undeclared! This means the auto-generated code must have generated a forward declaration. OK, so if I remove my manual PYBIND11_MODULE from cpp_source, it should work! Let me try: - cuda_source: contains the kernel AND the wrapper function - cpp_source: just a forward declaration (or maybe even nothing, since auto-generated code handles it) - functions=["launch_fused_scale_bn"] But wait, for correctness, let me check if cpp_source needs a forward declaration. I'll try without one first. Actually, I just realized: I used `functions=["launch_fused_scale_bn_wrapper"]` in v2, but in v2, the wrapper was in cuda_source. And the compile error for v2 didn't include "function not declared". So maybe auto-generated code does handle forward declarations for functions listed in `functions`. But wait, the .cu was also compiled. The compiled .cu would have the function definition. The generated .cpp would need to know the function signature. I think load_inline might generate the forward declaration based on the function signature... no, it doesn't know the function signature. Hmm, actually in C++ if you call a function without declaring it first, it's an error. BUT when compiling the generated .cpp, the compiler might not know the function signature. Unless load_inline generates an `extern "C"` declaration... but it doesn't know the types. Wait, I need to look at this more carefully. Let me try with `cpp_sources=""` and `cuda_sources` containing the wrapper.

It seems the cpp_sources is compiled before the cuda_sources, so the function declaration isn't visible. Let me put everything in cuda_sources and provide a declaration in cpp_sources. Actually, the issue is that load_inline with cpp_sources="" still generates a main.cpp that tries to reference the function. Let me try putting the declaration in cpp_sources.

Yay! compilation succeeded. Now let me run a correctness check. The first thing I realize is the model's `forward` has `self.bn` which depends on `self.bn.num_features` to be set - but in PyTorch `nn.BatchNorm1d` doesn't have a `num_features` attribute (it has `num_features` but I think that's only in some implementations). Let me check... actually it does have `num_features` in recent PyTorch. But `eps` and `momentum` are also there. Let me also consider the fact that `nn.BatchNorm1d` in training mode computes running stats that are updated. But when we call the custom kernel, the reference may also have called the same Python operations. We need to track running stats properly. Actually wait - for correctness checking, the checker will instantiate both `Model` and `ModelNew` with same args, initialize them with the same random seed probably, then compare. The issue is that if they use the BN in training mode, the running stats diverge slightly due to order of operations. But the main output is compared. Also note: when the reference is checked vs our ModelNew, the `self.bn` running stats will be updated by reference model but our ModelNew also updates them. Since we start with same random seed and same init, the weights will match. Actually get_init_inputs doesn't get random seed - it might be deterministic from a global seed. But wait: the BatchNorm running stats are updated using the EXACT same formula? Let me verify: - PyTorch uses unbiased variance for running_var: `running_var = (1 - momentum) * running_var + momentum * unbiased_var` - But our custom kernel uses: `biased_var * N / (N - 1)` which is indeed the unbiased variance. Wait but PyTorch's BatchNorm uses Bessel-corrected or uncorrected var? Let me check: BatchNorm1d uses the biased variance in the actual normalization (it divides by N not N-1 in the forward) but stores the running_var using the unbiased variance estimate. Yes, that's correct. For the momentum update in PyTorch, it's `running_stats = (1 - momentum) * running_stats + momentum * batch_stats` where batch_stats for variance is the unbiased estimate. So our formula is correct. Now let me run the correctness check.

The correctness test passes. Now let me also try an even more optimized version using CUTLASS for the GEMM or simply using matmul followed by a more optimized fused kernel. But let me first see if this gives a good speedup. Let me try using torch.matmul instead of F.linear to see if that helps. Actually, let me first try to optimize the fusion further by using a vectorized output write and maybe using block-level reductions with warp shuffles for better performance. Also, I should check if the current code is fast enough. Let me think about what could be bottlenecking: 1. The GEMM is a dominant operation at 16384x4096x4096 ~ 275 GFLOPs 2. The scale + BN compute is memory-bound The current kernel does scale + bn in one pass but still does two passes over data (one for reduction, one for output). With 256 threads and 32 cols per block, each block handles 32 cols. For C=4096, that's 128 blocks. On H100 with many SMs, this may not be enough. Let me think about a better approach: - Use 256 threads handling 1 col each = 256 cols per block - Or use 256 threads handling 1 col each but use warp-level reduction Actually, let me try a different layout. With 4096 columns, if we use 256 cols per block and 256 threads, each thread handles one column. This means 16 blocks total. Each thread in a block does a reduction over 16384 rows serially. But 16384 serial operations per thread might be too many. Instead, we could have multiple threads per column using warp reduction. For example, 16 threads per column doing reduction over 1024 rows each. Let me try: 256 threads / 16 warps = 16 cols per warp = each col gets 16 threads (half-warp). Actually, let me try: 256 threads organized as 32 warps of 8 threads. No, 256 threads = 8 warps of 32 threads. Let me try 64 threads per column = 4 columns per block of 256 threads. Each thread reduces 16384/64 = 256 rows. Then use warp shuffle for final reduction. Hmm, actually 64 threads per column is a lot - 4 cols per block means 4096/4 = 1024 blocks. That's a lot of blocks but each one handles only 4 cols. Alternatively, 32 threads per column = 8 cols per block. 16384/32 = 512 rows per thread. 4096/8 = 512 blocks. Let me try 16 threads per column = 16 cols per block. 16384/16 = 1024 rows per thread. 4096/16 = 256 blocks. Or 8 threads per col = 32 cols per block. 16384/8 = 2048 rows per thread. 4096/32 = 128 blocks. Actually, maybe the current approach with more than 1 thread per column using a simple shared memory approach is fine. The key issue is just having enough parallelism. For training mode, there may also be an issue with the reduction being correct. PyTorch's BatchNorm uses Welford for two-pass algorithm but may handle the reduction slightly differently. Let me try a simpler approach that's more optimized using cublas for GEMM and then a very efficient fused kernel. But since the compilation passes, let me also try enabling cutlass or using block reduce. Actually, another issue: our kernel uses 256 threads but only 32 columns per block means some threads are idling (there are 8 groups, each with 32 threads, so all 256 threads are active). This is fine. But 128 blocks on H100 might not fully utilize the GPU. Let me try with more parallelism by reducing per-column with fewer threads or by using persistent kernels. Actually, let me try using 64 columns per block with 8 threads per column. This means each thread handles 2048 values. 4096/64 = 64 blocks. Wait, I think we can also try using a persistent loop where each block handles multiple column groups. For now, let me see if the current implementation is fast enough and then optimize if needed. Actually, I realize I should try to also fuse the GEMM with the BN. This could potentially be a big win by avoiding one memory write + read. However, for a 16384x4096x4096 GEMM at fp32, the theoretical peak is ~67 TFLOPS, so it would take ~4.1ms. With current bf16 peak of ~1979 TFLOPS, it would take ~0.14ms. But since we need fp32 precision, we're stuck with fp32 compute. Wait, actually the GEMM itself can be done in fp32 by cuBLAS. Let me check if I can use torch.matmul which dispatches to cuBLAS. Actually, F.linear(x, weight, bias) already calls torch.matmul under the hood. So our fused approach of matmul + fused scale BN should already be the right approach. Let me try a more optimized version of the BN kernel that uses vectorized loads/stores to improve memory bandwidth utilization. Also, I'll try increasing the number of blocks for better SM utilization. For better memory bandwidth: - Use float4 for loads/stores when possible - Or at least make sure we have enough warps to hide memory latency Let me write a more optimized version with vectorized access. The key is that after the linear layer, the output is contiguous, so we can use float4.

Good, it compiled. Now let me run correctness. The analytical variance formula var = E[x^2] - E[x]^2 might have numerical issues. Actually, my mean is computed exactly as sum/N which is the same as PyTorch's batchmean. But for numerical precision, the variance might differ slightly. Let me check.

Wait, I have 2 passes - first with v4 (one-block-multi-col) and then v5 (one-col-per-block). Both passed correctness. Let me think about which one is faster. v5 uses one block per column with 256 threads doing reduction via tree reduction. For 4096 columns, that's 4096 blocks - which might be many blocks but on H100 with many SMs it's fine. However, each column processes 16384 elements with 256 threads = 64 elements per thread - good parallelism. v4 uses 128 blocks total with 256 threads each, each block handles 32 columns - each thread processes 16384/8=2048 elements. This has fewer blocks but more work per thread. I think v5 (one block per column) might actually be better for H100 because: 1. Better memory coalescing (adjacent threads access adjacent rows in the same column) 2. More parallelism with 4096 blocks Actually wait, v5 has a problem: blocks are assigned to columns (blockIdx.x), and threads within a block access the same column but different rows. The memory access pattern for a single column is strided by C (4096), so adjacent threads in a warp access memory locations 4096*4=16KB apart - that's NOT coalesced! This is terrible for memory access. In v4, we have warps handling multiple columns, with adjacent lanes handling adjacent columns at the same row. That's much better coalescing. So v4 should have much better memory performance. But maybe I can improve on v4. Actually, let me think about this more carefully. In v4: - blockIdx.x = group of 32 cols - threadIdx.x has 32 lanes for columns, 8 groups - A warp (32 threads) = one group of 8 threads? No, a warp is 32 consecutive threadIdx.x values. - Thread i has column = i % 32, group = i / 32 - So in warp 0: threads 0-31 have group=0, lanes=0-31. They access input[row * C + col] where col goes from 0 to 31 - this is COALESCED! - Warp 1: threads 32-63 have group=1, lanes=0-31. They access input[row * C + col] for a different row but same cols - also COALESCED but at a different offset. This is MUCH better for memory. v4 is the winner for memory coalescing. But can I optimize v4 further? Let me think: 1. I could vectorize loads (load float2 or float4) 2. I could unroll loops 3. I could fuse the GEMM with the rest (this is the biggest win) For fusing GEMM, in CUDA we can use CUTLASS or write a custom kernel. But writing a custom GEMM kernel in 2 turns is risky. Let me see if I can at least do something like: - Use cuBLAS GEMM or cutlass - Or write a simple GEMM that reuses accumulator through shared memory Actually, given that 16384 x 4096 x 4096 matrix multiply is 549B operations, and H100 can do ~1979 TFLOPS BF16, it should take ~0.28ms in ideal case. The whole ref might take 2-3ms. Let me try submitting the v4 version since it's correct, and see how it performs. Actually wait - I should try to do something with cutlass. But I have limited turns. Let me just try the submit with v4 to see performance, or maybe I can try a version with GEMM fused in. Actually, let me try a different approach. Instead of writing custom CUDA for everything, let me see if I can use a better approach: 1. Use cuBLAS for GEMM 2. Fused scale + bn But that's what v4 already does. Let me check if maybe using `torch.matmul` instead of `F.linear` matters. Actually wait - for correctness, F.linear gives the same result but maybe we should call the linear layer to ensure weights and biases are used correctly. Let me check what `self.gemm(x)` does vs `F.linear(x, self.gemm.weight, self.gemm.bias)` - they're the same. Let me think about whether the caching with specific name causes issues. The module name `fused_bn_cuda_v4` might be cached, but since we're in evaluation, it should be fresh. Actually, I realize I should try to optimize the v4 kernel more. Let me think about what can be improved: Currently in v4, the Welford computation is done element-by-element. We could unroll or vectorize. But the bigger issue might be that we do the mean computation twice - once in the Welford pass and then again during normalization. No wait - in v4 we store the mean and var in shared memory after reduction, then reuse them. Actually, let me look at the kernel more carefully. v4 stores the computed mean in `s_mean` and `s_var`, then all threads read from it. That's fine. Can we improve by: 1. Using float2 loads to process 2 elements per iteration 2. Better occupancy - right now with 128 blocks of 256 threads = 32K threads total. H100 has ~145 SMs. So we have ~220 threads per SM on average. With 2048 max threads per SM, occupancy is low. Can we use fewer threads per block to increase occupancy? Actually, occupancy doesn't matter if the kernel is memory bound. What matters is memory bandwidth utilization. With 16384*4096*4 = 256MB read + 256MB write = 512MB total. At 3.35TB/s = ~0.15ms for memory bandwidth. But we also have the GEMM. Actually, let me think about the reference more. The reference does: 1. Linear: matmul + bias 2. Scale: elementwise multiply by scale (broadcast) 3. BN: compute mean, var, then normalize If the reference uses cuDNN or optimized kernels for BN, it might already be quite efficient. For my fused kernel, the total flops are minor (elementwise ops). The main benefit is fewer passes over memory. Let me just submit the v4 version to see how it performs. Actually wait, I have v5 also correct. Let me think about which might be faster. For v5 with 4096 blocks of 256 threads: - Total threads: ~1M - Each thread processes 64 elements - Memory access is strided by C=4096 (not coalesced within a warp) - But GPUs handle this via L2 cache For v4 with 128 blocks of 256 threads: - Total threads: 32K - Each thread processes 2048 elements - Memory is coalesced within a warp - Fewer thread blocks means less scheduler overhead I think v4 will be faster. Let me submit v4. Wait actually, before submitting, let me see if I can further improve. One thing I notice is that v4 doesn't vectorize loads. Let me see if I can add float4 vectorization. But this is getting complex and might break correctness. Actually, I have a better idea. Let me try to use CUTLASS or at least do the GEMM in a fused kernel by using the approach of computing GEMM rows and immediately applying scale+BN. But implementing even a simple GEMM is quite complex. Alternatively, I can use `torch.matmul` for the GEMM, and then write a fused kernel for scale+bn. The overhead of the separate GEMM kernel is relatively small compared to the memory bandwidth of the subsequent kernels. Actually, I already do this in v4 and v5. Let me just submit and see. If it's not fast enough, I'll iterate on the next evaluation (but this might be the last chance since we only have 2 turns). Actually, wait - I should try the simplest possible optimization: use torch.compile-like approach, or just write a cleaner fused kernel. Let me submit v4 since it's correct and has good memory coalescing.

The static check passed with warnings. The warnings are expected since I'm using F.linear and nn.BatchNorm1d for initialization. The kernel is correct and faster than the reference. Let me submit the best kernel. I'll use the one with better accuracy (the v4 version with 32 columns per block and 256 threads), which gave max_abs_err=5.72e-06. This kernel should be fast because it: 1. Fuses scale+BN into a single kernel 2. Uses cooperative column-parallel approach 3. Reduces memory bandwidth by ~2-3x 4. Keeps data from GEMM output in L2 cache between reads