We need to write a custom CUDA kernel that implements Layer Normalization for the given shape: input is 4D tensor (batch_size=16, features=64, dim1=256, dim2=256). The LayerNorm is over normalized_shape = (features, dim1, dim2). In PyTorch nn.LayerNorm with normalized_shape = (features, dim1, dim2), it normalizes over the last 3 dims (i.e., for each batch element, compute mean and variance over the last 3 dimensions). The parameters gamma (weight) and beta (bias) are of shape (features, dim1, dim2). But they are optional: default is elementwise_affine=True, so weight and bias are created of shape normalized_shape. Yes, the reference uses default elementwise_affine=True. So weight shape = (64,256,256); bias shape = (64,256,256). The input size per batch element is huge: 64*256*256 = 4,194,304 elements. There are 16 batch elements. The total number of elements = 67,108,864. That's a lot. We need a faster CUDA kernel than PyTorch's LayerNorm. We need to write a custom kernel. Since normalized_shape is large (4M elements per instance), we need high-bandwidth, compute-optimized kernel. LayerNorm formula: For each batch element (i.e., each outermost index), compute mean and variance over all elements in the last 3 dims. Then normalize: y = (x - mean) * rsqrt(var + eps) * weight + bias. Here var is unbiased? In PyTorch LayerNorm, var is computed as the mean of squared deviations (i.e., population variance), but with unbiased=False? Wait: PyTorch's LayerNorm computes mean and variance over the normalized dimensions. It uses torch.nn.functional.layer_norm, which computes variance with unbiased=False (Bessel's correction is not applied). Let's verify: The PyTorch implementation computes mean, then var = mean((x - mean)^2). Yes, it's population variance. Then y = (x - mean) / sqrt(var + eps) * weight + bias. We can implement a custom kernel that fuses mean/var computation and the affine transform. Because weight and bias are per-element across the normalized dims. Potential strategies: - Use cooperative groups or warp shuffle to reduce within blocks. - Because per-instance size is ~4M elements, we can process each batch element with multiple blocks. For each batch element, compute mean and variance using a two-pass or one-pass (Welford) algorithm across the elements. Since we need mean and variance before writing output, we have to either: - Do a two-pass approach: first pass compute mean and var, store in shared memory or global memory, then second pass apply normalization. But we can't launch separate kernels easily within load_inline wrapper (though we can call multiple functions, but can also do a grid-level reduction). Alternatively, we can have each block compute partial sum and partial sum of squares, write to global memory, then do a second kernel that reads partials and computes final mean/var and writes output. However, using two separate functions means we need to call two CUDA functions from Python, but overhead likely okay since it's just two launches. But we could also use a single kernel with two phases using cooperative groups (grid sync), but that's tricky without cooperative launch. Simpler: do two-pass within the same kernel per block by processing each batch element within a block? But each batch element is huge; a block cannot cover all elements. We could use a single thread block per batch element? But block size max is 1024 threads, which is tiny compared to 4M elements. We need many blocks per batch element. But if we have many blocks per batch element, they can't share data easily unless we use atomic operations or global scratch memory. Alternative: Use a two-kernel approach. Kernel 1: compute per-block partial sums and partial sum of squares. Kernel 2: reduce the partials to final mean and var, then apply normalization. Let's design: Input x: shape [B, C, H, W] where B=16, C=64, H=256, W=256. N = C*H*W. Number of blocks per batch element: we could assign each block to process a contiguous chunk of elements. For sum reduction, each block processes e.g., 1024 * 4 = 4096 elements per block (each thread loads 4 elements). So total blocks per batch = N / (block_size * elems_per_thread). If N = 4,194,304, block_size=256, elems_per_thread=4 => block_elems=1024 => blocks per batch = 4096. So total blocks = B * blocks_per_batch = 16 * 4096 = 65536 blocks. That's a lot but okay. Kernel 1: Each block computes local sum and local sum_sq for its chunk. It writes these to global arrays of shape [B, num_blocks_per_batch, 2]. Or two separate arrays: partial_sums [B, num_blocks_per_batch], partial_sq [B, num_blocks_per_batch]. Kernel 2: Each block (or each thread) processes one batch element. It reads all partial sums for that batch element and reduces them to final mean and var. Then each block processes a chunk of the output: it loads x, weight, bias, computes normalized output, writes it. However, we still need to parallelize the write across many blocks. But kernel 2 can be configured such that each block processes a chunk of elements for a given batch element. For example, we can have gridDim.x = B * num_blocks_per_batch, where blockIdx.x maps to (b, block_id). But each block needs the global mean and var for its batch element. We can precompute mean and var per batch element in shared memory or load them from a small device array [B]. But each block would need to know mean[b] and inv_std[b]. We can compute them in kernel 2: before processing its chunk, each block reads mean and inv_std from global memory. But we need to ensure that the reduction of partials to mean/inv_std is done before any block uses them. One approach: launch a small kernel that does the reduction per batch element and writes mean and inv_std to arrays [B]. Then launch a third kernel that applies the affine transform using those arrays. That would be three kernels. Or we can combine reduction and transform in a single kernel by having each block do its own reduction of partials? But partials are many (4096 per batch). Each block could read all partials for its batch element? If there are 4096 partials per batch, and each block reads 4096 floats, that's 16 KB, which fits in shared memory. But if we have many blocks per batch element, each block would redundantly read the same partials and compute the same mean and var. That's okay! Redundancy is fine as long as it's not too heavy. If we have 4096 blocks per batch element, each block reads 4096 partials => 4096 * 4096 = 16M reads per batch element, which might be expensive. But we could assign a dedicated block (or warp) to compute mean/inv_std and store in shared/global, then have other blocks wait? Without grid sync, hard. Better: 3-kernel approach: Kernel 1: Compute per-block partial sum and sum_sq. Write to global arrays. Kernel 2: For each batch element, one thread block (or a few threads) reduces the partials to compute mean and inv_std (rsqrt(var + eps)) and writes to arrays mean[B], inv_std[B]. Kernel 3: Each block processes a chunk of elements, reads x, weight, bias, mean[b], inv_std[b], computes y, writes output. This is straightforward and likely fast enough. PyTorch probably does something similar but maybe with more optimization. However, we can also optimize kernel 3 by vectorized loads (float4) to saturate memory bandwidth. The weight and bias are also large (4M each). So memory bandwidth is the bottleneck. Using float4 loads can help. Let's think about GPU specs. Need to get GPU specs. Let's call get_gpu_specs.

The task is to write a custom CUDA kernel for LayerNorm. Let me analyze the problem: 1. The input shape is `(batch_size, features, dim1, dim2)` = `(16, 64, 256, 256)` 2. The normalized_shape is `(features, dim1, dim2)` = `(64, 256, 256)` 3. So we're normalizing over the last 3 dimensions: features × dim1 × dim2 = 64 × 256 × 256 = 4,194,304 elements per batch item This is a very large normalization dimension. With batch size 16, total elements = 16 × 4,194,304 = 67,108,864. For LayerNorm, we need to: 1. Compute mean over the last 3 dims for each batch item 2. Compute variance over the last 3 dims for each batch item 3. Normalize: `(x - mean) / sqrt(var + eps) * weight + bias` Since the normalization dimension is huge (4M elements), this is memory bandwidth bound. The key optimization is to fuse everything into a single kernel and use efficient reductions. For a reduction over 4M elements with 16 batch items, we need tree-based reduction in shared memory. Let me think about the kernel design: - Each batch item gets processed independently - For each batch item, we need to reduce 4M elements - We can launch 1 block per batch item, or multiple blocks per batch item - With 4M elements and block size 256, that's ~16384 elements per thread for a single block approach Actually, for 4M elements, a single block with 256 threads would need each thread to process 16K elements, which is too many. We should use multiple blocks. But wait - with 16 batch items, we only need 16 independent reductions. Using one block per batch item with 1024 threads: - 4,194,304 / 1024 = 4096 elements per thread - That's still a lot but manageable. Or we could use a two-pass approach: - First pass: each block computes partial sum/sum-sq, writes to global memory - Second pass: single block reads partial results and finalizes For 4M elements with block size 256: - We'd need 16384 blocks per batch item, which is too many for a 2-pass if we use global atomics Let me think more carefully. With batch_size=16, we can process each batch item with cooperative groups or use grid-stride loops. Actually, for layer norm with such large dims, the standard approach is: 1. First pass: compute mean and variance using warp/block reduction 2. Second pass: normalize But computing mean and var in one go is better using Welford's algorithm. Given the huge reduction dimension (4M), let's do: - Grid: (num_batches, 1, 1) where num_batches = 16 - Block: 256 or 512 or 1024 threads - Each thread processes N elements via grid-stride / block-stride loop Wait, if we use 1 block per batch with 1024 threads, each thread handles 4096 elements. A block-stride loop of 4096 iterations is fine. The reduction can be done with shared memory. Actually, we can also use multiple blocks per batch item for better parallelism. With 16 SMs and batch_size=16, 1 block per batch gives us 1 block per SM, which could under-utilize the GPU if each block takes a long time. But 4096 iterations × 16 batches = pretty good utilization. Let me think about occupancy. With 1024 threads per block, that's 32 warps. On H100, maximum warps per SM is 64 (I think - actually for Hopper it's 64 warps/SM = 2048 threads). So 1 block of 1024 threads gives 50% occupancy per SM. With 16 blocks across 132 SMs (H100 has 132 SMs I think? Let me check - actually H100 SXM has 132 SMs, but the 80GB version might have slightly fewer, around 132-144). Wait, the spec didn't give SM count. H100 typically has 132 or 144 SMs. With 16 blocks of 1024 threads, even with 132 SMs we only use 16 SMs, leaving most idle. That's bad! So we should use MORE blocks per batch item. Let's use something like: - Block size: 256 threads - Blocks per batch: enough to fill the GPU - With 132 SMs and wanting 2 blocks per SM for occupancy, we want ~264 blocks total - With 16 batches, that's ~16-17 blocks per batch But each block needs to handle partial reduction and we need a second pass to combine. Two-pass approach: - Pass 1: Each block computes partial mean and partial m2 (for variance) for its chunk of elements - We store partial results in global memory (2 floats per block) - Pass 2: A single block per batch reads all partial results and computes final mean and var - Pass 3: All blocks normalize their elements But 3 passes is a lot of global memory traffic. Can we do better? Alternative: Use atomic operations to accumulate mean and m2. But atomics on floats are slow and non-deterministic. Another alternative: Tree reduction using multiple blocks with cooperative groups. CUDA has cg::grid_group for grid-sync. But grid sync is tricky and can deadlock. Let me reconsider: With block size 256 and 16 blocks per batch: - 16 × 16 = 256 blocks total - Each thread handles 4,194,304 / (16 × 256) = 1024 elements For 2-pass approach: - Pass 1: 256 blocks, each computes local mean and m2, stores to workspace - 16 batches × 16 blocks = 256 partial sums - Pass 2: 16 blocks (1 per batch), each reads 16 partial results, computes final mean and var, stores to workspace - Pass 3: 256 blocks normalize elements The workspace needs: - Pass 1 output: 16 batches × 16 blocks × 2 floats = 512 floats = 2KB - Pass 2 output: 16 batches × 2 floats = 128 floats Memory traffic: - Read input once in pass 1 - Read input again in pass 3 - Write output in pass 3 - Plus small workspace traffic Total read: 2 × input = 2 × 16 × 64 × 256 × 256 × 4 bytes = 512 MB Total write: 1 × output = 256 MB Total: 768 MB At 3.35 TB/s, theoretical minimum = 768 / 3350 = 0.23 ms. That's very fast. But kernel launch overhead for 3 passes might be significant. Can we fuse pass 2 and pass 3? Actually, we can use persistent blocks with a 2-kernel approach: 1. Compute mean and var (2-pass) 2. Normalize (1-pass) Or even better, for such a large problem, maybe we should use one of these strategies: **Strategy A: Single kernel with sequential access** - Launch 256 blocks total (16 per batch) - Each block processes elements and does a block-level reduction - Then we need cross-block reduction... Actually, let's think about this differently. For LayerNorm with such large dimensions, PyTorch likely already uses highly optimized cuDNN/cuBLAS kernels. To beat it, we need to be very efficient. The most practical approach for beating PyTorch's LayerNorm on large dims: 1. Use vectorized loads (float4) to maximize memory bandwidth 2. Fuse mean/var computation and normalization into as few passes as possible 3. Use efficient warp-shuffle reductions Let me look at what typical custom LayerNorm kernels do. For large N, a common pattern is: - Kernel 1: Each block computes partial sums using shared memory, write to global - Kernel 2: Final reduction + normalization Or we can use a single kernel with multi-stage approach: - Stage 1: All blocks compute partial statistics - __syncthreads() at grid level (cooperative groups) - Stage 2: One block finalizes statistics - Stage 3: All blocks normalize But grid-level sync is complex. Let me try a 2-kernel approach with optimal parallelism: **Kernel 1: Compute mean and variance** - Grid: Many blocks per batch - Each block: uses sequential float4 loads, accumulates sum and sum-of-squares - Block-level reduction via shared memory + warp shuffle - AtomicAdd or workspace-based accumulation Wait, if we use atomics for the cross-block reduction within each batch, we can do it in one kernel for stats + one kernel for normalization. But atomics on float32 across many blocks can be slow. Alternative: Let each block write its partial stats. Use a second small kernel to reduce partial stats. Then a third kernel for normalization. Actually for 16 batches with ~16-128 blocks each, the partial stats array is tiny, so a second kernel is very fast. The third kernel is where most time is spent. 2-kernel approach is better than 3-kernel: - Kernel 1: Compute partial stats + normalize in same kernel? No, we need final stats first. What if we do: - Kernel 1: Blocks compute partial stats and store. When a designated block (e.g., block 0) for each batch sees all partials are done via atomic counter, it computes final mean/var and stores. - Then all blocks normalize. This is a single kernel using atomic counter. But we'd need to re-read the data. Hmm, let me just go with the clean 2-kernel approach: - Kernel 1 (stats): Grid with many blocks per batch. Each thread loads and accumulates. Block reduces to shared memory. Block writes partial mean and var to global workspace. Then a small grid reduces the partials per batch and writes final mean and var per batch. - Actually we can run this as two kernels or fuse the partial reduction. Let's just do: Kernel 1a for partial stats, Kernel 1b for final stats reduction, Kernel 2 for normalization. Wait, Kernel 1b is tiny (16 items × maybe 128 partials = 2048 items), so it's negligible. We can also just do it on CPU since the output is 16 × 2 = 128 floats. That's probably fastest! So: 1. CUDA Kernel: partial_sum_kernel - launches blocks per batch, each outputs partial count/sum/sum_sq 2. CPU: reduce partials to final mean and var for each batch 3. CUDA Kernel: normalize_kernel - uses final mean and var from CPU Actually, nuance: reduction from GPU to CPU requires a sync. Then we need to launch the normalize kernel. That's 2 round trips but still might be fast for such large data. Wait, copying a few floats from GPU to CPU and back is ~1-2 microseconds of latency. For a kernel that takes hundreds of microseconds, that's fine. But there's an even better approach: Use online algorithm and shared memory to do block-level reduction, then use a single reduction block per batch. Let me design the kernels more concretely. **Input**: [16, 64, 256, 256] = [batch, C, H, W] **Normalized shape**: [64, 256, 256] = 4,194,304 elements **Output**: same shape Let's use `float4` loads. Total elements per batch = 4,194,304. With float4, that's 1,048,576 float4s. If block_size = 256, and we use 64 blocks per batch: - Total blocks = 16 × 64 = 1024 blocks - Elements per block = 4,194,304 / 64 = 65,536 float4s... wait that doesn't divide evenly. Actually 4,194,304 / 64 = 65,536. And each block with 256 threads processes 65,536 / 256 = 256 float4s per thread. That's a reasonable loop count. Or let's use 128 blocks per batch: 4,194,304 / 128 = 32,768 elements per block. With 256 threads, each handles 128 float4s. Total blocks = 2048. Let's go with 128 blocks per batch, block_size = 256. **Kernel 1: compute_stats** - Each thread loads 128 float4s (512 floats) - Accumulates sum and sum of squares (or uses Welford for numerical stability) - Block reduces sum and sum_sq to shared memory - Thread 0 writes block-level sum and sum_sq to global workspace - Grid: (16, 128) or linear with 2048 blocks Wait, we need to be careful about grid dimensions for simplicity. Let's use a 1D grid. Grid size = batch_size * blocks_per_batch = 16 * 128 = 2048 Block size = 256 Each block knows its batch index: batch_idx = blockIdx.x / blocks_per_batch Within-batch block index: block_in_batch = blockIdx.x % blocks_per_batch Each thread in a block processes elements starting from: base = batch_idx * N + block_in_batch * block_elements + threadIdx.x With block-stride = blockDim.x, process all elements in the block's chunk. After reduction, write partial_sum[batch_idx * blocks_per_batch + block_in_batch] and partial_sq[...]. **CPU reduction**: For each batch, sum all partial_sums and partial_sqs, compute mean = sum / N, var = sq / N - mean^2 (or m2 / N). **Kernel 2: normalize** - Grid same as Kernel 1: 2048 blocks - Each thread loads same elements - Applies normalization: out = (x - mean) * rsqrt(var + eps) * weight + bias - Writes using float4 stores However, the weights and biases are only for the normalized shape. Wait, let me re-read the PyTorch LayerNorm. In PyTorch's LayerNorm, if normalized_shape = (C, H, W), then weight and bias have shape (C, H, W) as well. They're learnable parameters. Wait, the reference doesn't specify elementwise_affine=True, but nn.LayerNorm defaults to elementwise_affine=True, so weight and bias ARE used and have shape equal to normalized_shape. So we need to apply: `y = (x - mean) * rstd * weight + bias` where weight[i, j, k] and bias[i, j, k] match the last 3 dims. Since the weight and bias have shape [64, 256, 256] and our data is also [16, 64, 256, 256], the weight/bias are the same across batch dimension. For vectorized loads, we need to load weight and bias too. We can load them as float4 and broadcast the batch dimension. Wait, but the weights and biases are parameters stored in the nn.LayerNorm module. In ModelNew, we need to either: 1. Copy them to the kernel at each call 2. Or store them in our ModelNew and pass them The most efficient approach: In ModelNew.__init__, store the normalized_shape. In forward, if ln is the layer norm module, we'll get weight and bias, or create ones/zeros if they're None. For the CUDA kernel, we need to pass weight and bias as tensors. Actually wait, the problem says `ModelNew` should accept the same constructor arguments. So it gets `normalized_shape`. We should probably create our own weight and bias parameters, or we can simply call the reference ln's forward for parameters that don't need CUDA optimization. But for correctness and speed, we should: 1. Create weight and bias in __init__ (same as nn.LayerNorm does) 2. In forward, compute mean/rstd with CUDA, then apply affine transform Or even simpler: Just use the nn.LayerNorm module for weight/bias and only replace the actual computation with CUDA kernel... but then we're calling nn.LayerNorm which might have overhead. Actually, the task says "Replace the forward computation... with a custom CUDA kernel." So we need to implement forward() ourselves. Let me think about the best approach. A very efficient approach for large layer norm: ```python class ModelNew(nn.Module): def __init__(self, normalized_shape): super().__init__() self.normalized_shape = normalized_shape self.weight = nn.Parameter(torch.ones(normalized_shape)) self.bias = nn.Parameter(torch.zeros(normalized_shape)) self.eps = 1e-5 def forward(self, x): # Launch custom CUDA kernel return custom_ln(x, self.weight, self.bias, self.eps) ``` In the CUDA kernel, for each batch item, we compute mean and rstd over all elements, then normalize with weight and bias. Given that the normalized dimension is 4M elements, let me think about whether to use warp-level primitives or shared memory for reduction. For 128 blocks per batch, 256 threads per block: - Each thread loop: 4096 elements / 256 threads = wait, 4,194,304 / 128 blocks = 32,768 elements per block. With 256 threads, each thread handles 128 elements. - Using float4: 32,768 / 256 = 128 float4 per thread. That's a good amount. For reduction within a block: - Use warp shuffle to reduce to lane 0 - Then lane 0 writes to shared memory - Thread 0 reduces shared memory values Should be very efficient. Let me write the kernel code. For numerical stability, Welford's algorithm is better for variance, but for 4M elements with random [0,1] inputs, simple sum and sum_sq should be fine. Actually let me use Welford for better accuracy. Wait, Welford per thread then merge across threads... that's complex. For simplicity, I'll just use double precision for accumulation. The input is fp32, but accumulating in fp64 virtually eliminates numerical issues. H100 has great fp64 performance (34 TFLOPS). Actually, since we want to be fast and fp64 is only 34 TFLOPS vs 67 TFLOPS for fp32, and this is bandwidth bound anyway, using fp64 for accumulation is fine. The accumulator is register-bound. Let me use double (or float for speed if we can get away with it). For fp32 inputs in [0,1] with 4M elements, the mean is ~0.5 and variance ~0.08. The sum will be ~2M which fits exactly in fp32 (fp32 has 24 bits of mantissa = ~16M exact integers). Sum of squares ~1.3M. Actually sum of 4M values each ~0.5 is ~2M, which is well within fp32 precision. Sum of squares of 4M values each ~0.33 is ~1.3M. So fp32 accumulation should be exact to about 1 part in 2^24 / 2M ≈ 8. That might not be enough. Actually for correctness tolerance of 1e-4 against reference, fp32 accumulation might be okay. But let's be safe and use double for the stats accumulation. Wait, PyTorch LayerNorm also uses some accumulation strategy. Let me check: PyTorch uses Welford's with fp32 or mixed precision. For fp32 inputs, it typically accumulates in fp32 with Welford. So we might be fine with fp32 + Welford. Actually, no - for the tolerance check, we just need to match PyTorch's output. PyTorch on CUDA might use cuDNN's LayerNorm which probably uses something sophisticated. But I think a simple mean/variance with fp64 accumulation or Welford with fp32 will be close enough. Hmm, actually for numerical equivalence, the simplest approach is to compute mean and variance the same way PyTorch does. But we don't have access to cuDNN internals. Let me just use Welford's online algorithm with float accumulation, which is numerically stable and should match or be very close to standard implementations. Or simpler: compute mean first, then variance in a second pass. But that requires reading the data twice in the stats kernel. A 2-pass variance computation is actually very accurate. Wait, but we want to minimize memory bandwidth. Is there a single-pass approach that's accurate enough? For [0,1] random data with 4M elements: - Mean ≈ 0.5, variance ≈ 1/12 ≈ 0.083 - The two-pass algorithm: first compute mean, then compute E[(x-mean)^2] Two-pass in kernel: 1. First loop: accumulate sum 2. Block reduce to get mean 3. Second loop: accumulate (x - mean)^2 4. Block reduce to get variance But that requires iterating over the data twice within each thread! For 128 iterations per thread, that's still fine. Actually we can do even better. We only need one pass over memory for stats if we use Welford. But Welford with warp reduction is more complex to implement. Let me think about what's simplest and fastest: - Option A: Two-pass within block. Each thread loops twice. Total thread work: 256 loops. - Option B: Welford single-pass. Each thread loops once. Total thread work: 128 loops, but more ops per element. For bandwidth-bound kernels, the loop count matters less than memory bandwidth. Option A might actually be fine. But wait - we need to do this for 128 blocks per batch, and then we also have the normalize kernel which reads memory again. So total reads = 2 for stats + 1 for normalize = 3 passes over memory. Actually let me reconsider. Can we read memory once for stats and once for normalize? Yes, 2 total. But for two-pass stats, we need 2 passes just for stats, so 3 total. For a 2-kernel approach: - Kernel 1: Each block does partial sum. Needs 1 pass. Writes partial sums. - Small CPU/GPU reduction for final mean and var. - Kernel 2: Normalize. Needs 1 pass (read input, read weight, write output). Total: 2 passes over input. Much better! For the variance, we can use the "one pass with partial stats" approach too. Each block computes: 1. sum of x 2. sum of x^2 Then var = E[x^2] - (E[x])^2 This is the traditional two-moment approach and is less numerically stable than Welford. For our data, though, it should be fine. But let's use a correction: m2 = sum_sq - mean*sum = sum((x-mean)^2) computed in a second pass... no wait. Actually for maximum accuracy matching PyTorch, let me just use their approach. PyTorch's native layer norm on CUDA uses something like Welford or two-pass. Let me look at PyTorch's implementation... For CUDA, PyTorch uses ATen/native/cuda/layer_norm_kernel.cu which typically does: 1. Compute mean 2. Compute rstd = rsqrt(var + eps) where var is computed as mean of (x - mean)^2 In their kernel, they often use a single kernel that does both for small normalized sizes, or multiple passes for large sizes. For large sizes, cuDNN might handle it. For our custom kernel, let's do: 1. Compute mean per batch with a reduction kernel 2. Compute variance per batch with another reduction kernel (using the mean) 3. Normalize with a third kernel That's 3 kernels, 3 passes over input. But each kernel is simple. Or even simpler, just fuse steps 2 and 3: compute variance and normalize in one kernel. That gives 2 passes total. Actually, the very best approach for beating PyTorch might be: - Kernel 1: Compute mean (reduction) - Kernel 2: Compute rstd and output normalized values Wait, can we do better? For 4M elements, if we use a single block per batch for reduction, that's 1 block doing 4M elements with 1024 threads = 4096 elements per thread. With 16 batches, that's only 16 blocks total. On H100 with 132 SMs, that's terrible occupancy! So we need multiple blocks per batch. Let's use 256 blocks per batch with 256 threads each. That gives 4096 blocks total, which is great for occupancy. With 256 blocks per batch: - Elements per block = 4,194,304 / 256 = 16,384 = 64 float4s per thread Total: 16 * 256 = 4096 blocks. Good utilization. Wait, 4,194,304 / 256 = 16,384. With 256 threads per block, each thread handles 64 elements. With float4, that's 16 float4 per thread. Very reasonable. **Kernel 1 (mean_reduction)**: - Grid: 4096 blocks - Block: 256 threads - Each thread loads 16 float4s (64 floats) and sums them - Block reduces sum to thread 0 - Thread 0 atomically adds to global mean[batch_idx] Actually, atomicAdd on float across 256 blocks per batch. That's 256 atomics to the same address per batch. This might be slow due to contention, but with 16 batches hitting 16 different addresses, it could be okay. But 256 atomics to one address is a lot. Better to use workspace + separate reduction. Let me use the workspace approach: - Grid 4096 blocks - Each block writes partial sum to workspace[batch_idx * blocks_per_batch + block_in_batch] - Then a small kernel with 16 blocks (1 per batch) and 256 threads reduces 256 partial sums Actually, since 16 × 256 = 4096 partial sums is small, we can even launch a kernel with 16 blocks where each block loads 256 floats and reduces them. Or use a single block with 1024 threads to reduce all 4096 values. Or even simpler: the CPU can do this reduction since we only have 4096 floats! We can just do `partial_sums.reshape(16, 256).sum(dim=1)` in PyTorch, which is very fast. Then copy to GPU or just keep on GPU. Actually, partial_sums is already on GPU. Using torch.sum on GPU is fine and very fast for such a small tensor. So: 1. Kernel 1: compute partial sums per block, store in workspace 2. `mean = partial_sums.view(batch_size, -1).sum(dim=1) / N` (GPU or CPU) 3. Kernel 2: compute partial sum of (x - mean)^2 per block, store in workspace 4. `var = partial_vars.view(batch_size, -1).sum(dim=1) / N` (GPU or CPU) 5. `rstd = 1.0 / sqrt(var + eps)` (GPU or CPU) 6. Kernel 3: normalize using mean and rstd Steps 2, 4, 5 are tiny and can be done on GPU with a single small kernel or just PyTorch ops. But this is 3 passes over input (kernel 1, kernel 2, kernel 3). That's a lot of memory bandwidth. Can we combine kernel 1 and kernel 2? If we compute mean first, then in kernel 2 we can both compute variance AND normalize? No, we need all variances before normalizing any. Wait, what if we compute mean in kernel 1. Then in kernel 2, each block computes its local (x - mean)^2 and normalizes its elements immediately... but we need the global variance for normalization. So no. Unless we do something clever: since var = E[(x - mean)^2], and each block computes its local sum of squared deviations. The global var is the average of these. For the normalization, we need `1/sqrt(global_var + eps)`. If each block computes local sum of sq dev, writes to global, then we can't normalize in the same kernel. So minimum is 2 passes over input if we can somehow get stats in 1 pass... or maybe 2 kernels with clever interleaving. Actually, for maximum performance, the cleanest approach that still has a chance of beating PyTorch: - Compute mean and rstd using **one** reduction kernel (Welford's algorithm, block-level reduction, write partial stats) - Reduce partial stats to final mean and rstd (small kernel or PyTorch) - Normalize (one kernel) This is 2 passes over input + tiny overhead for partial stat reduction. Much better than 3. Let me design the single-pass Welford kernel. **Welford's algorithm** per thread: ```cpp float mean = 0.0f; float m2 = 0.0f; int count = 0; for each x: count++; float delta = x - mean; mean += delta / count; float delta2 = x - mean; m2 += delta * delta2; ``` After processing all elements, each thread has (count, mean, m2). Then we need to merge Welford states across threads. The merge formula for two states (n1, mean1, m2_1) and (n2, mean2, m2_2): ```cpp float delta = mean2 - mean1; float total_count = n1 + n2; float new_mean = (n1 * mean1 + n2 * mean2) / total_count; float new_m2 = m2_1 + m2_2 + delta * delta * n1 * n2 / total_count; ``` We can do this with warp shuffle, then shared memory. After block-level reduction, each block writes (mean, m2) to workspace. Then we reduce across blocks using the same merge. Since we process complete chunks with equal count per thread, the count is the same for all threads, so the merge can be simplified. But let's keep it general. Actually, for a simpler implementation, since all threads in a block process the same number of elements (or differ by at most 1 if N doesn't divide evenly), we could just use sum and sum-of-squared-deviations-from-mean. But we don't know the mean upfront. Wait, simpler: We can compute sum and sum_sq in one pass. Then var = sum_sq/N - mean^2. For numerical stability, we can run a second loop to compute sum of (x - mean)^2... but that's 2 loops. For single loop, let's just use Kahan summation or accept the slight numerical difference. Since tolerance is 1e-4 and inputs are fp32, the naive variance should be fine. Actually, let me check: For random numbers in [0,1], mean=0.5, E[x^2] = 1/3 ≈ 0.3333, var = 1/12 ≈ 0.08333. sum ≈ 2,097,152 sum_sq ≈ 1,398,101 If we compute mean = sum/N = 2,097,152/4,194,304 = 0.5 exactly if the random generator is uniform. With real random, there will be some error. For var = sum_sq/N - mean^2: sum_sq/N = 1,398,101 / 4,194,304 = 0.333333... mean^2 = 0.25 The subtraction 0.3333 - 0.25 = 0.08333. In fp32, this should be accurate to about 1e-6 relative, or about 1e-7 absolute. Well within our 1e-4 tolerance. But wait - the reference might use Welford, which gives slightly different results. The key is whether the difference exceeds 1e-4. For the scale of these values, it shouldn't. Let me just go with sum and sum_sq in double precision to be safe on numerical accuracy. Then compute var as sum_sq/N - mean*mean. This will be extremely accurate. **Updated kernel design:** **Kernel 1: partial_stats_kernel** - Grid: batch_size * blocks_per_batch blocks - Block: 256 threads - Each thread loops over its elements, loading with float4 - Accumulates sum and sum_sq in double - Block reduces sum and sum_sq using warp-shuffle then shared memory - Thread 0 writes block-level stats to workspace **CPU/GPU reduction**: - For each batch, sum all partial sums and partial sum_sqs - Compute mean = total_sum / N - Compute var = total_sum_sq / N - mean * mean - Compute rstd = rsqrt(var + eps) **Kernel 2: normalize_kernel** - Grid: same as kernel 1 - Each thread loops, loads input with float4, loads weight and bias with float4 - Normalizes: out = (x - mean) * rstd * weight + bias - Stores with float4 Wait, weight and bias have shape [64, 256, 256]. To load them with float4, we need to be careful about indexing. The input is [B, C, H, W]. For a given global thread index, we need to map to the correct (b, c, h, w) and then index the weight as [c, h, w]. For the normalize kernel with float4: - Input is treated as float array of size B*C*H*W - Each float4 load reads 4 consecutive floats - We need weight and bias loaded correspondingly For weight at position [c, h, w*4 : w*4+4], we can load as float4 if w is aligned to 4. Given H=256, W=256, W is divisible by 4, so float4 alignment is great! For a global index `idx` into the flattened [B, C, H, W] array: - b = idx / (C*H*W) - rem = idx % (C*H*W) - c = rem / (H*W) - hw = rem % (H*W) - h = hw / W - w = hw % W For float4 at `idx` (which must be aligned to 4): - We load input[idx:idx+4] - The corresponding weight elements are at [c, h, w: w+4] in the weight tensor Since W=256 is divisible by 4, the float4 chunk [c, h, w:w+4] is contiguous in memory (Tensor is stored row-major as [C, H, W]). Wait, PyTorch tensor layout: For a tensor of shape [C, H, W], the strides are [H*W, W, 1]. So elements [c, h, w], [c, h, w+1], [c, h, w+2], [c, h, w+3] are contiguous. As long as w % 4 == 0, we can load them as float4. For the flattened input of shape [B, C, H, W], strides are [C*H*W, H*W, W, 1]. So four consecutive elements in the flattened array could span: - Within same (b, c, h, w:w+3): contiguous, good - Across h dimension boundary: if w=253, we'd get w=253,254,255, then next is h+1, w=0. These are contiguous in memory! Because the next element after [c, h, 255] is [c, h+1, 0] if H*W layout is used. Wait no: For [C, H, W] with strides [H*W, W, 1]: - [c, h, 255] is at offset c*H*W + h*W + 255 - [c, h+1, 0] is at offset c*H*W + (h+1)*W + 0 = c*H*W + h*W + W + 0 If W=256, then [c, h, 255] offset +1 = c*H*W + h*W + 256 = [c, h+1, 0]. Yes! So the entire [C, H, W] tensor is just a contiguous C*H*W array. Similarly for [B, C, H, W], it's a contiguous B*C*H*W array. So if we treat the input as a contiguous 1D array and use float4 loads at positions 0, 4, 8, ..., the float4s always contain valid contiguous elements. For weight of shape [C, H, W] = [64, 256, 256], it's also contiguous. At position offset = c*H*W + h*W + w, loading float4 gives [c, h, w:w+4]. So the indexing is straightforward! Now, for the normalize kernel, each thread needs to know: - batch_idx (for mean and rstd lookup) - linear offset into the tensor (for input, output, weight, bias) The mean and rstd are per-batch, stored as [B] arrays. Let me write the CUDA code now. Wait, but we need to handle the case where N is not divisible by 4 for the float4 loads. N = 4,194,304 which IS divisible by 4 (and by 256). Good! For partial_stats_kernel: ```cpp extern "C" __global__ void partial_stats_kernel( const float* __restrict__ input, double* __restrict__ partial_sums, double* __restrict__ partial_sum_sqs, int batch_size, int N) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; int tid = threadIdx.x; int total_threads_per_batch = blocks_per_batch * blockDim.x; int global_tid = block_in_batch * blockDim.x + tid; const float* batch_input = input + batch_idx * N; double sum = 0.0; double sum_sq = 0.0; // Grid-stride loop for (int i = global_tid; i < N / 4; i += total_threads_per_batch / 4) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; sum += val.x; sum_sq += val.x * val.x; sum += val.y; sum_sq += val.y * val.y; sum += val.z; sum_sq += val.z * val.z; sum += val.w; sum_sq += val.w * val.w; } // Handle remaining elements if N not divisible by 4 // But N is divisible by 4 in our case. // Block reduction __shared__ double s_sum[256]; __shared__ double s_sum_sq[256]; s_sum[tid] = sum; s_sum_sq[tid] = sum_sq; __syncthreads(); // Reduction in shared memory for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { s_sum[tid] += s_sum[tid + s]; s_sum_sq[tid] += s_sum_sq[tid + s]; } __syncthreads(); } if (tid == 0) { int idx = batch_idx * blocks_per_batch + block_in_batch; partial_sums[idx] = s_sum[0]; partial_sum_sqs[idx] = s_sum_sq[0]; } } ``` Wait, there's an issue. If we do `N / 4` float4 loads, but `total_threads_per_batch / 4` doesn't evenly divide, we might have issues. Let me think more carefully. Actually, `total_threads_per_batch` should be a multiple of... wait, no. In the grid-stride loop, `i` goes from `global_tid` to `N/4` in steps of `total_threads_per_batch`. But `global_tid` is an index into float4 elements. So `total_threads_per_batch` should be the number of float4-processing threads. If blockDim.x = 256 and blocks_per_batch = 128: - total_threads_per_batch = 128 * 256 = 32768 threads - Each thread processes N/4 / 32768 = 1,048,576 / 32768 = 32 float4s Wait, N = 4,194,304. N/4 = 1,048,576 float4s. 1,048,576 / 32,768 = 32 exactly. So each thread processes exactly 32 float4s. Very clean! Let me recalculate with blocks_per_batch = 64: - total_threads = 64 * 256 = 16,384 - 1,048,576 / 16,384 = 64 float4s per thread With blocks_per_batch = 256: - total_threads = 256 * 256 = 65,536 - 1,048,576 / 65,536 = 16 float4s per thread I think 128 blocks per batch with 32 float4s per thread is good. Let me use 64 blocks per batch = 64 float4s per thread. Or 256 blocks = 16 float4s per thread. For better SM utilization, let's use 256 blocks per batch. Total blocks = 16 * 256 = 4096. - Each thread: 16 float4 loads = 64 elements - Block size: 256 threads This should provide excellent parallelism on H100. Actually, let me reconsider block size. 256 threads is standard. 4096 blocks of 256 threads = 1,048,576 threads. H100 can handle this easily. I'll use `blocks_per_batch = 256`. But this is a constant. Or we can pass it as a parameter... Actually, since our grid size needs to match blocks_per_batch, we need to decide at compile time or pass it. To keep it simple, let me define `BLOCKS_PER_BATCH` as a constant. Since N = 4,194,304 and N/4 = 1,048,576, we want 1,048,576 to be divisible by (BLOCKS_PER_BATCH * blockDim.x). 1,048,576 = 2^20. So any combination where BLOCKS_PER_BATCH * blockDim.x is a power of 2 ≤ 2^20 works. blockDim.x = 256 = 2^8. BLOCKS_PER_BATCH = 256 = 2^8. Total = 2^16 = 65,536. 1,048,576 / 65,536 = 16. Or BLOCKS_PER_BATCH = 128 = 2^7. Total = 2^15 = 32,768. 1,048,576 / 32,768 = 32. Or BLOCKS_PER_BATCH = 64. Total = 16,384. Quotient = 64. I'll use BLOCKS_PER_BATCH = 128 so each thread handles 32 float4s. That's a nice middle ground. Wait, what if the kernel is used for different sizes? The problem says ModelNew should accept the same constructor arguments as Model, so it needs to work for any normalized_shape. But to be safe, I should write the kernel generically, not hardcoded for (64, 256, 256). For a generic kernel, we can pass N as an argument and use dynamic grid sizing. Generic approach: - Compute total_elements_per_batch = product of normalized_shape - Compute num_blocks_per_batch based on total_elements - Or just fix block size and compute blocks_per_batch by formula Let me use a generic formula: ```cpp int elements_per_batch = N; // total elements per batch int threads_per_block = 256; int blocks_per_batch = max(1, min((N + threads_per_block * 32 - 1) / (threads_per_block * 32), 1024)); ``` Wait, we want blocks_per_batch to be a power of 2 or something nice for reduction? Not necessarily if we use PyTorch for the final reduction. Actually, for the generic case, we don't know N. Let's just use: ```cpp int total_threads = blocks_per_batch * threads_per_block; ``` Where blocks_per_batch is computed on host side and passed to kernel. For the loop: ```cpp for (int i = global_tid; i < N / 4; i += total_threads) { // load float4 } ``` And for the tail: ```cpp int remainder_start = (N / 4) * 4; for (int i = global_tid + remainder_start; i < N; i += total_threads) { // load single float } ``` This works generically. Now for the normalize kernel, same structure but with affine transform. For weight and bias, we need to handle the case where they're None (elementwise_affine=False). In PyTorch LayerNorm, elementwise_affine defaults to True but can be False. For the normalize kernel: ```cpp extern "C" __global__ void normalize_kernel( const float* __restrict__ input, const float* __restrict__ weight, const float* __restrict__ bias, float* __restrict__ output, const float* __restrict__ mean, const float* __restrict__ rstd, int batch_size, int N, int blocks_per_batch) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; int tid = threadIdx.x; int total_threads_per_batch = blocks_per_batch * blockDim.x; int global_tid = block_in_batch * blockDim.x + tid; float m = mean[batch_idx]; float rs = rstd[batch_idx]; const float* batch_input = input + batch_idx * N; float* batch_output = output + batch_idx * N; // Process with float4 for (int i = global_tid; i < N / 4; i += total_threads_per_batch) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; float4 w = (weight != nullptr) ? reinterpret_cast<const float4*>(weight)[i % (N / 4)] : make_float4(1.0f, 1.0f, 1.0f, 1.0f); float4 b = (bias != nullptr) ? reinterpret_cast<const float4*>(bias)[i % (N / 4)] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); val.x = (val.x - m) * rs * w.x + b.x; val.y = (val.y - m) * rs * w.y + b.y; val.z = (val.z - m) * rs * w.z + b.z; val.w = (val.w - m) * rs * w.w + b.w; reinterpret_cast<float4*>(batch_output)[i] = val; } // Handle tail int remainder_start = (N / 4) * 4; for (int i = global_tid + remainder_start; i < N; i += total_threads_per_batch) { int w_idx = i % N; // Actually for weight, it's just i since weight has shape [N] when flattened float w = (weight != nullptr) ? weight[i % N] : 1.0f; float b = (bias != nullptr) ? bias[i % N] : 0.0f; batch_output[i] = (batch_input[i] - m) * rs * w + b; } } ``` Wait, there's a subtle issue with weight indexing. When we treat input as [B, N] flattened, weight has shape [N]. So weight[i] for the i-th element in the normalized dimensions. Since we compute `i` as a global offset into the [B, N] tensor, and `i % N` gives the position within the normalized shape. Actually, wait. For the float4 case, if weight is shape [N], then weight float4 at position `i` corresponds to elements `4*i, 4*i+1, 4*i+2, 4*i+3`. Since the input and weight have the same layout, we can use the same `i` index for both. But in our loop, `i` ranges over `0` to `N/4 - 1`. For batch 0, the input float4 elements are at `i = 0, 1, 2, ...`. For batch 1, the input float4 elements would start at `N/4`. But since `batch_input = input + batch_idx * N`, the float4 pointer is `reinterpret_cast<const float4*>(batch_input)`, so indexing at `i` gives the i-th float4 in that batch's data. That's correct. For weight, we don't offset by batch. So the float4 pointer for weight is just `reinterpret_cast<const float4*>(weight)`, and for input position `i` within the batch, the corresponding weight position is also `i` (since weight has no batch dimension). Wait, but if batch_input starts at `input + batch_idx * N`, and it's cast to `float4*`, then `batch_input[i]` is `*(float4*)(input + batch_idx * N + i * 4)`. This is correct. For weight, `weight[i]` is `*(float4*)(weight + i * 4)`. And since weight is [C, H, W] = [64, 256, 256], it's contiguous. The elements `weight[i*4 : i*4+4]` correspond to the same logical positions as `batch_input[i*4 : i*4+4] % N`. Is this true? For batch_input at `batch_idx * N + i * 4`, the position within the normalized shape is `(i * 4) % N`. And the weight pointer `weight + i * 4` points to `weight[0, 0, i*4]` if we think of weight as 1D... wait, that's not right. If weight is shape [C, H, W] and flattened to 1D, then weight[k] corresponds to a specific (c, h, w). For the batch_input at position `batch_idx * N + k`, it corresponds to the same (c, h, w). So `batch_input[k]` maps to `weight[k % N]`. When we do `reinterpret_cast<const float4*>(batch_input)[i]`, this gives elements `batch_input[batch_idx * N + i * 4]`, `batch_input[batch_idx * N + i * 4 + 1]`, etc. For weight, we need `weight[(i * 4) % N]`, `weight[(i * 4 + 1) % N]`, etc. Since `i * 4 < N` (because `i < N/4`), `(i * 4) % N = i * 4`. So `weight[i * 4]` through `weight[i * 4 + 3]` are exactly `reinterpret_cast<const float4*>(weight)[i]`. Yes! Since `i < N/4`, we never wrap around. So `weight[i]` as float4* is correct. BUT: what if N is not divisible by 4? Then `N/4 * 4 < N`, and the last few elements don't form a complete float4. We handle this with the tail loop. For the tail loop, `i` goes from `N/4 * 4` to `N-1`. For each such `i`, `weight[i % N] = weight[i]` since `i < N`. Correct. So the indexing works! One more issue: `make_float4` might not be available. We can use literal initialization or `float4{1.0f, 1.0f, 1.0f, 1.0f}` in C++. Actually in CUDA C++, we can use: ```cpp float4 w_val; if (weight != nullptr) { w_val = reinterpret_cast<const float4*>(weight)[i]; } else { w_val = make_float4(1.0f, 1.0f, 1.0f, 1.0f); } ``` But `make_float4` is available in CUDA headers. To be safe, I'll use it. Now let's think about the launch parameters. For the stats kernel: - threads_per_block: 256 - blocks_per_batch: Let's calculate dynamically. We want each thread to process a reasonable amount of data, maybe 64-256 floats (16-64 float4s). - blocks_per_batch = max(1, N / (threads_per_block * 64)) or similar. - But we don't want too few blocks (bad occupancy) or too many (too many partials). For N = 4,194,304: If target_elements_per_thread = 64 = 16 float4s: blocks_per_batch = N / (256 * 64) = 4,194,304 / 16,384 = 256. So 256 blocks per batch, 16 float4s per thread. Total blocks = 4096. Actually, let me calculate this more carefully and pick good parameters. H100 has ~132 SMs. Each SM can run multiple blocks. With our kernel, each block uses some registers and shared memory. Let me estimate. For the stats kernel: - Each thread: ~10 registers for loop variables + 2 doubles (16 registers) = maybe 26 32-bit registers? Actually double variables count as 2 registers each in 32-bit register file. So maybe 30-40 registers per thread. - 256 threads: ~10K registers per block. Well under 64K limit. - Shared memory: 256 * sizeof(double) * 2 = 4KB per block. - Blocks per SM: limited by 64K registers / 10K = ~6 blocks. And shared memory: 227KB / 4KB = ~56 blocks. So register limited to ~6 blocks per SM. 132 SMs * 6 = 792 blocks concurrently. With total 4096 blocks, that's fine. For occupancy, we want independent warps to hide latency. With 256 threads per block = 8 warps, and we're doing sequential loads, warps within a block are executing together. But with 4096 blocks, there's plenty of work for all SMs. Actually, a potential issue: In the stats kernel, each thread in a block cooperates via shared memory for reduction. Blocks are independent. So 4096 blocks works well. Wait, there's another consideration. The grid-stride loop in each thread processes elements spaced by total_threads_per_batch. For good memory coalescing, adjacent threads should access adjacent memory. In our loop: ```cpp for (int i = global_tid; i < N / 4; i += total_threads_per_batch) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; } ``` If `global_tid` is computed as `block_in_batch * blockDim.x + tid`, then in block 0, threads 0,1,2,... access float4 elements 0,1,2,... which are contiguous. Perfect coalescing! In block 1, threads access elements starting at 256, 257, etc. Also contiguous and coalesced. Good. But what about L2 cache? Different blocks access different regions. Since we might have 256 blocks per batch, and N/4 = 1,048,576 float4s, each block covers 1,048,576 / 256 = 4,096 float4s = 16KB. This is small and fits in L2. Actually, I'm overthinking this. The memory access pattern is linear and coalesced, which is optimal. One potential optimization: We might want to unroll the loop. But the compiler should handle this for constant-length loops. Since 32 or 16 iterations per thread is small, the compiler might already unroll it. Another optimization: Use `__ldg` for read-only data. Since we only read input once (for stats) and once (for normalize), `__ldg` can help by using the texture/readonly cache. Also: For the stats kernel, we might want to use float instead of double if we can get away with it for speed. Double operations are ~2x slower in throughput. But since we're bandwidth bound, the difference might be small. Let me keep double for accuracy and see how it performs. Wait, another idea: instead of computing sum_sq in double, what if we use compensated summation (Kahan) in float? That would be faster but still accurate. For simplicity, let me just use double. Actually, for even better performance, we could avoid the 3-kernel/2-kernel overhead and do everything in one kernel using atomic counters for a grid sync. But that's complex and error-prone. Let me stick with 2 kernels. Let me write the complete code now. Wait, one more thing. We need to handle N not being a multiple of total_threads_per_batch. For the grid-stride loop, this is handled naturally. But for coalescing, if N/4 is not divisible by total_threads_per_batch, the last iteration of some threads will have fewer elements or wrap around? No, the condition `i < N/4` handles it. For the case where N is not divisible by 4, the tail loop handles remaining 0-3 elements. Actually, what if N < total_threads_per_batch? Then many threads won't enter the loop. That's fine. Now, writing the Python file: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline cuda_source = r""" #include <cuda_runtime.h> #include <math.h> extern "C" __global__ void partial_stats_kernel( const float* __restrict__ input, double* __restrict__ partial_sums, double* __restrict__ partial_sum_sqs, int batch_size, int N, int blocks_per_batch) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; int tid = threadIdx.x; int total_threads_per_batch = blocks_per_batch * blockDim.x; int global_tid = block_in_batch * blockDim.x + tid; if (batch_idx >= batch_size) return; const float* batch_input = input + batch_idx * (size_t)N; double sum = 0.0; double sum_sq = 0.0; // Process float4 chunks int num_float4 = N / 4; for (int i = global_tid; i < num_float4; i += total_threads_per_batch) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; sum += val.x; sum_sq += val.x * val.x; sum += val.y; sum_sq += val.y * val.y; sum += val.z; sum_sq += val.z * val.z; sum += val.w; sum_sq += val.w * val.w; } // Handle remaining elements int remainder_start = num_float4 * 4; for (int i = global_tid + remainder_start; i < N; i += total_threads_per_batch) { float val = batch_input[i]; sum += val; sum_sq += val * val; } __shared__ double s_sum[256]; __shared__ double s_sum_sq[256]; s_sum[tid] = sum; s_sum_sq[tid] = sum_sq; __syncthreads(); // Block reduction #pragma unroll for (int s = 128; s > 0; s >>= 1) { if (tid < s) { s_sum[tid] += s_sum[tid + s]; s_sum_sq[tid] += s_sum_sq[tid + s]; } __syncthreads(); } if (tid == 0) { int idx = batch_idx * blocks_per_batch + block_in_batch; partial_sums[idx] = s_sum[0]; partial_sum_sqs[idx] = s_sum_sq[0]; } } extern "C" __global__ void normalize_kernel( const float* __restrict__ input, const float* __restrict__ weight, const float* __restrict__ bias, float* __restrict__ output, const float* __restrict__ mean, const float* __restrict__ rstd, int batch_size, int N, int blocks_per_batch) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; int tid = threadIdx.x; int total_threads_per_batch = blocks_per_batch * blockDim.x; int global_tid = block_in_batch * blockDim.x + tid; if (batch_idx >= batch_size) return; float m = mean[batch_idx]; float rs = rstd[batch_idx]; const float* batch_input = input + batch_idx * (size_t)N; float* batch_output = output + batch_idx * (size_t)N; bool has_weight = (weight != nullptr); bool has_bias = (bias != nullptr); // Process float4 chunks int num_float4 = N / 4; for (int i = global_tid; i < num_float4; i += total_threads_per_batch) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; float4 w = has_weight ? reinterpret_cast<const float4*>(weight)[i] : make_float4(1.0f, 1.0f, 1.0f, 1.0f); float4 b = has_bias ? reinterpret_cast<const float4*>(bias)[i] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); val.x = (val.x - m) * rs * w.x + b.x; val.y = (val.y - m) * rs * w.y + b.y; val.z = (val.z - m) * rs * w.z + b.z; val.w = (val.w - m) * rs * w.w + b.w; reinterpret_cast<float4*>(batch_output)[i] = val; } // Handle remaining elements int remainder_start = num_float4 * 4; for (int i = global_tid + remainder_start; i < N; i += total_threads_per_batch) { float w = has_weight ? weight[i] : 1.0f; float b = has_bias ? bias[i] : 0.0f; batch_output[i] = (batch_input[i] - m) * rs * w + b; } } """ _cpp_source = "" # Not needed since we use extern "C" functions # Compile the inline CUDA extension custom_ln = load_inline( name="custom_ln", cpp_sources=_cpp_source, cuda_sources=cuda_source, functions=["partial_stats_kernel", "normalize_kernel"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False, ) ``` Wait, I need to be careful with `load_inline`: - `cpp_sources` can be empty or contain includes - `cuda_sources` contains the CUDA code - `functions` lists the names of the functions to expose - The functions should be declared with `extern "C"` to avoid C++ name mangling Also, I need to make sure the CUDA code compiles. Let me think about potential issues: 1. `reinterpret_cast` on `__restrict__` pointers - this should be fine. 2. `size_t` arithmetic for large N. 3. The reduction with `s = 128` assumes blockDim.x = 256. For other block sizes, we'd need a different reduction. Let me make it dynamic or just assume blockDim.x = 256. But for safety, maybe I should use `blockDim.x / 2` and a loop. Actually, using a compile-time constant for warp reduction is dangerous if blockDim.x ever changes. But in this code, we control blockDim.x, so it's fine. I'll make it dynamic for safety though, or we can check that blockDim.x <= 256. Wait, if blockDim.x > 256, the `s_sum[256]` array is too small. Let me just set blockDim.x = 256 and use shared array of size 256. We can assert or just always launch with 256. Actually, the simplest approach is to use shared memory sized by `extern __shared__` or just use warp shuffle reduction instead of shared memory to avoid shared memory size issues. Let me use warp shuffle reduction instead of shared memory for the final steps. It's simpler and doesn't need shared memory allocation. Warp shuffle reduction: ```cpp // Within-warp reduction for (int offset = warpSize / 2; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xffffffff, sum, offset); sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset); } __shared__ double s_sum[32]; // One per warp __shared__ double s_sum_sq[32]; int lane = tid % warpSize; int warp_id = tid / warpSize; if (lane == 0) { s_sum[warp_id] = sum; s_sum_sq[warp_id] = sum_sq; } __syncthreads(); // Block leader reduces warp results if (tid < blockDim.x / warpSize) { sum = s_sum[lane]; sum_sq = s_sum_sq[lane]; for (int offset = (blockDim.x / warpSize) / 2; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xffffffff, sum, offset); sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset); } if (lane == 0) { s_sum[0] = sum; s_sum_sq[0] = sum_sq; } } __syncthreads(); if (tid == 0) { // write results } ``` This is more code but doesn't depend on blockDim.x being a specific value. But since we're launching with 256 threads, it's fine either way. Actually, for simplicity, I'll just use shared memory reduction but with a dynamic size up to 256. I'll use `extern __shared__ double s_data[];` and allocate dynamically. But for `load_inline`, dynamic shared memory can be tricky because we need to pass the size in the kernel launch... actually no, we can just use `extern __shared__` and specify size at launch with the third argument <<<grid, block, shared_mem>>>. Let me just hardcode block_size = 256 to keep things simple. The code is cleaner that way. Actually, to be safe and generic, let me use warp shuffle. Here's the updated reduction: ```cpp // Warp-level reduction #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); } __shared__ double s_sum[32]; __shared__ double s_sum_sq[32]; int lane = tid & 31; int warp_id = tid >> 5; if (lane == 0) { s_sum[warp_id] = sum; s_sum_sq[warp_id] = sum_sq; } __syncthreads(); // Final reduction by warp 0 if (warp_id == 0) { sum = (lane < blockDim.x / 32) ? s_sum[lane] : 0.0; sum_sq = (lane < blockDim.x / 32) ? s_sum_sq[lane] : 0.0; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); } if (lane == 0) { s_sum[0] = sum; s_sum_sq[0] = sum_sq; } } __syncthreads(); if (tid == 0) { int idx = batch_idx * blocks_per_batch + block_in_batch; partial_sums[idx] = s_sum[0]; partial_sum_sqs[idx] = s_sum_sq[0]; } ``` This is robust for any block size that's a multiple of 32 and <= 1024. Now, for the PyTorch side: ```python class ModelNew(nn.Module): def __init__(self, normalized_shape): super().__init__() self.normalized_shape = tuple(normalized_shape) if not isinstance(normalized_shape, tuple) else normalized_shape self.eps = 1e-5 # Create weight and bias like PyTorch's LayerNorm self.weight = nn.Parameter(torch.ones(self.normalized_shape)) self.bias = nn.Parameter(torch.zeros(self.normalized_shape)) def forward(self, x): batch_size = x.shape[0] N = x[0].numel() # elements per batch item # Reshape input to [batch_size, N] x_view = x.reshape(batch_size, N) # Determine blocks per batch threads_per_block = 256 elements_per_thread = 64 # target blocks_per_batch = max(1, (N + threads_per_block * elements_per_thread - 1) // (threads_per_block * elements_per_thread)) # Ensure it's not too large blocks_per_batch = min(blocks_per_batch, 1024) # For power-of-2 N like 4194304, we'd get blocks_per_batch = 256 total_blocks = batch_size * blocks_per_batch # Workspace for partial sums partial_sums = torch.empty(batch_size * blocks_per_batch, dtype=torch.float64, device=x.device) partial_sum_sqs = torch.empty(batch_size * blocks_per_batch, dtype=torch.float64, device=x.device) # Launch stats kernel custom_ln.partial_stats_kernel( x_view, partial_sums, partial_sum_sqs, batch_size, N, blocks_per_batch, grid=(total_blocks,), block=(threads_per_block,) ) # Reduce partial stats partial_sums_view = partial_sums.view(batch_size, blocks_per_batch) partial_sum_sqs_view = partial_sum_sqs.view(batch_size, blocks_per_batch) total_sum = partial_sums_view.sum(dim=1) total_sum_sq = partial_sum_sqs_view.sum(dim=1) mean = total_sum / N var = total_sum_sq / N - mean * mean rstd = torch.rsqrt(var + self.eps).float() mean = mean.float() # Launch normalize kernel output = torch.empty_like(x_view) w = self.weight.view(N) if self.weight is not None else None b = self.bias.view(N) if self.bias is not None else None custom_ln.normalize_kernel( x_view, w, b, output, mean, rstd, batch_size, N, blocks_per_batch, grid=(total_blocks,), block=(threads_per_block,) ) return output.view_as(x) ``` Wait, I need to think about the performance of the PyTorch reduction step: ```python total_sum = partial_sums_view.sum(dim=1) total_sum_sq = partial_sum_sqs_view.sum(dim=1) ``` For batch_size=16 and blocks_per_batch=256, this is summing [16, 256] tensors. This is very fast (a few microseconds on GPU). But it requires launching a small CUDA kernel via PyTorch. The issue is that this creates a synchronization point between the two custom kernels. But since the first custom kernel is the bottleneck, the overhead should be minimal. However, to maximize performance, we could compute the final stats with a custom CUDA kernel too. But for 16x256, PyTorch is probably faster. Another optimization: We can avoid creating `x_view` and just pass `x` to the kernel, computing the flattened indexing inside the kernel or relying on PyTorch's contiguous storage. Since x is [B, C, H, W] and contiguous (as created by torch.rand), we can pass `x` directly to the kernel and treat it as a flat array of size B*N. But the kernel signature expects `const float*` and a linear index. Since `x` is contiguous, `x.data_ptr()` works. Similarly for output, we can create `torch.empty_like(x)` and get its data ptr. Let me simplify: don't reshape to x_view. Just pass x directly. Actually, for the weight and bias, they have shape [C, H, W] = [64, 256, 256] and are contiguous. So `self.weight.data_ptr()` works with the kernel treating weight as a flat array of size N. Now, one issue: when we call `custom_ln.partial_stats_kernel(...)`, the load_inline generates a pybind11 binding that expects torch tensors. So we should pass tensors directly. Actually, load_inline with `functions=["partial_stats_kernel", "normalize_kernel"]` auto-generates wrappers that accept torch tensors and extract their data pointers. But I need to verify the exact API. Looking at the PyTorch docs for `load_inline`: the generated functions take torch tensors as arguments and launch the kernel. But I need to pass grid and block dimensions. The typical usage is: ```python custom_ln.partial_stats_kernel(arg1, arg2, ..., blocks=grid_size, threads=block_size) ``` Wait, actually `load_inline` generates functions that launch CUDA kernels with the syntax: ```python func(arg1, arg2, ..., grid=<tuple>, block=<tuple>) ``` Or maybe it's: ```python func.argtypes = [...] func(grid, block, arg1, arg2, ...) ``` Hmm, I'm not entirely sure. Let me think about how `load_inline` works. `torch.utils.cpp_extension.load_inline` compiles the code and returns a module. The functions in `functions=[...]` are made callable. For CUDA kernels, the wrapper is auto-generated. Looking at PyTorch source code, the generated wrapper uses pybind11 and for CUDA functions, it might extract tensor pointers and call the kernel with `<<<grid, block>>>`. Actually, from memory and looking at common patterns: ```python custom_op.func(tensor1, tensor2, grid=(gx, gy, gz), block=(bx, by, bz)) ``` But I'm not 100% certain. Let me think about an alternative: I can define a C++ wrapper function that takes at::Tensor, extracts data_ptr(), computes grid dimensions, and calls the CUDA kernel. Then list that wrapper in `functions=[...]`. Yes! This is much more reliable and gives us full control. ```cpp torch::Tensor layer_norm_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, float eps) { int batch_size = input.size(0); int N = input.numel() / batch_size; // compute mean and rstd // ... auto output = torch::empty_like(input); // launch kernel // ... return output; } ``` But wait, the instructions say "Do not write your own `PYBIND11_MODULE(...)` block in `cpp_sources`. `load_inline` auto-generates one from the `functions=[...]` argument." So I should write C++ functions that take torch tensors and return tensors, and declare them as `extern "C"`. But actually, if they're in the `cuda_sources`, they need to be C++ because we're using ATen/Torch types. Wait, `cuda_sources` is just CUDA C++ source that gets compiled with the PyTorch headers. I can write C++ wrapper functions in the CUDA source. But I also have `cpp_sources` for pure C++. Let me use `cpp_sources` for the wrapper function that calls the CUDA kernels. Actually, let me re-read the instructions: "List the host-side wrapper function names in `functions=[...]` instead." So I can write wrapper functions in either `cpp_sources` or `cuda_sources`, and they'll be exposed via pybind11. For a wrapper function: ```cpp #include <torch/extension.h> #include <cuda_runtime.h> // Forward declarations of CUDA kernels void partial_stats_kernel(...); void normalize_kernel(...); at::Tensor custom_layer_norm(at::Tensor input, at::Tensor weight, at::Tensor bias, float eps) { // ... implementation return output; } ``` And `functions=["custom_layer_norm"]`. But these wrapper functions need access to the CUDA kernels. Since they're in the same compilation unit (or linked together), it works. In `load_inline`, both `cpp_sources` and `cuda_sources` are compiled and linked together. So I can declare the kernels as `extern "C"` in the CUDA source and call them from C++. Wait, `extern "C"` functions in CUDA source are decorated with `__global__`, which means they're CUDA kernels. Can normal C++ code call them? Yes, with the `<<< >>>` syntax. Actually, in C++ host code, calling a `__global__` function requires the `.cu` file or at least NVCC compilation. Since `cpp_sources` is compiled with the C++ compiler (not NVCC), it can't contain `<<< >>>` kernel launches! However, in `load_inline`, `cpp_sources` might still be compiled with nvcc if CUDA is available. Let me think... Actually, looking at `torch.utils.cpp_extension.load_inline` (https://pytorch.org/docs/stable/cpp_extension.html): it uses `torch.utils.cpp_extension.load` internally. The `cpp_sources` are compiled as C++ and `cuda_sources` as CUDA. Both are linked together. But C++ code in `cpp_sources` cannot directly launch kernels with `<<< >>>` unless it's in a `.cu` file. So either: 1. Put everything in `cuda_sources` (which is compiled as `.cu`) 2. Or declare a wrapper in `cuda_sources` that does the launch, and put any pure C++ helper in `cpp_sources` The simplest approach: put the wrapper function in `cuda_sources` as well. Since `cuda_sources` is compiled with nvcc, it can contain both host C++ and device CUDA code. So our `cuda_sources` would contain: 1. The `__global__` kernel definitions 2. A host C++ wrapper function that computes grid dims, allocates tensors, launches kernels, and returns the result And `cpp_sources` can be empty or contain includes. Wait, but the instructions say "Do not write your own PYBIND11_MODULE". The `functions=[...]` argument tells `load_inline` which functions to wrap. So our wrapper function just needs to be declared in the compiled code. Let me design the complete solution: ```cpp // In cuda_sources: #include <torch/extension.h> #include <cuda_runtime.h> #include <ATen/cuda/CUDAContext.h> // CUDA kernels... extern "C" __global__ void partial_stats_kernel(...) { ... } extern "C" __global__ void normalize_kernel(...) { ... } at::Tensor custom_layer_norm_forward( at::Tensor input, at::Tensor weight, at::Tensor bias, float eps) { int batch_size = input.size(0); int N = input.numel() / batch_size; // Allocation auto partial_sums = at::empty({batch_size, blocks_per_batch}, input.options().dtype(at::kDouble)); auto partial_sum_sqs = at::empty({batch_size, blocks_per_batch}, input.options().dtype(at::kDouble)); // Get stream auto stream = at::cuda::getCurrentCUDAStream(); // Launch stats kernel partial_stats_kernel<<<grid, block, 0, stream>>>(...); // Reduce using ATen auto total_sum = partial_sums.sum({1}); auto total_sum_sq = partial_sum_sqs.sum({1}); auto mean = total_sum / N; auto var = total_sum_sq / N - mean * mean; auto rstd = at::rsqrt(var + eps); // Output auto output = at::empty_like(input); // Launch normalize kernel normalize_kernel<<<grid, block, 0, stream>>>(...); return output; } ``` Wait, there's a problem: PyTorch's ATen sum on a CUDA tensor is asynchronous, but we're using its result immediately to compute `mean`, `var`, etc. These operations build a computational graph. Then we launch `normalize_kernel` using `rstd.data_ptr<float>()`. But does this synchronize properly? In PyTorch, `.data_ptr()` on a tensor waits for all ops on that tensor to complete. So when we access `rstd.data_ptr()`, PyTorch ensures the reduction and computation are done. But then we pass it to a raw CUDA kernel launch. The kernel launch is asynchronous, which is fine because it operates on the data that is now ready. However, there's a potential issue: if we do `auto stream = at::cuda::getCurrentCUDAStream()`, and launch our kernel on that stream, it's fine because the previous operations (reductions) are also on that stream (or synchronize implicitly when we access data_ptr). Actually, ATen operations might not all be on the same stream or might require explicit synchronization. For safety, we could use `mean = mean.cpu()` to force sync? No, that would be slow. Alternatively, we could just ensure that our kernels operate in the proper stream and that PyTorch's lazy execution ensures everything is ordered correctly. One important detail: When we do `at::rsqrt()`, it creates a new tensor. To get `float` pointers for our CUDA kernel, we need to ensure the tensors are contiguous on CUDA. Actually, in modern PyTorch, the sequence: ```cpp auto total_sum = partial_sums.sum({1}); auto mean = total_sum / N; ``` returns tensors on GPU, and when we do `mean.data_ptr<float>()`, it implicitly blocks until `mean` is computed. But we also need to make sure `mean` and `rstd` are in the right stream. Let me think about this more. CUDA kernels launched via `<<< >>>` in PyTorch C++ extensions typically use the default/current stream. The ATen operations also use the current stream. So operations should be implicitly ordered as long as they're on the same stream. But when we extract `data_ptr()` and then launch another kernel, there should be an implicit dependency because the second kernel reads from the memory that the first operations write to. CUDA runtime might not know this unless we're using Graphs or explicit events. Actually, when `mean.data_ptr()` is accessed, it might trigger a host-side synchronization if the tensor isn't ready. But I'm not 100% sure. In practice, many PyTorch C++ extensions do exactly this and it works. To be extra safe, we can do everything with pure CUDA kernels and avoid ATen in the middle. But that requires writing a reduction kernel for the small [batch, blocks] tensor, which is simple. Let me write a small reduction kernel for the partial stats: ```cpp extern "C" __global__ void reduce_stats_kernel( const double* __restrict__ partial_sums, const double* __restrict__ partial_sum_sqs, float* __restrict__ mean, float* __restrict__ rstd, int batch_size, int blocks_per_batch, float eps) { int batch_idx = blockIdx.x; if (batch_idx >= batch_size) return; double sum = 0.0; double sum_sq = 0.0; const double* batch_sums = partial_sums + batch_idx * blocks_per_batch; const double* batch_sum_sqs = partial_sum_sqs + batch_idx * blocks_per_batch; for (int i = threadIdx.x; i < blocks_per_batch; i += blockDim.x) { sum += batch_sums[i]; sum_sq += batch_sum_sqs[i]; } // Block reduction __shared__ double s_sum[256]; __shared__ double s_sum_sq[256]; s_sum[threadIdx.x] = sum; s_sum_sq[threadIdx.x] = sum_sq; __syncthreads(); for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (threadIdx.x < s) { s_sum[threadIdx.x] += s_sum[threadIdx.x + s]; s_sum_sq[threadIdx.x] += s_sum_sq[threadIdx.x + s]; } __syncthreads(); } if (threadIdx.x == 0) { double n = static_cast<double>(...); // Wait, we need N // Actually we don't have N here! We'd need to pass it. } } ``` Wait, if I do a custom reduction kernel, I'd need to pass `N` to compute mean and variance. That's easy. But then I'd have a tiny third kernel (16 blocks, 256 threads) which is fine. But is 3 kernels really slower than using PyTorch's reduction + potential sync issues? Maybe not by much. Let me just use custom kernels for everything to avoid synchronization issues and have maximum control. Actually, let me simplify the design: - Stats kernel: computes partial sums - Reduce kernel: computes final mean and rstd - Normalize kernel: applies normalization Stats kernel and normalize kernel are the big ones. Reduce kernel is trivial. Wait, can we combine stats and reduce? If we use atomics for sum and sum_sq per batch, we don't need a separate reduce kernel. Each block atomically adds its partial sum to `global_sum[batch_idx]`. Then a single block per batch computes mean and rstd. But we'd still need a way to signal when all blocks are done (atomic counter) before computing rstd. That gets us back to grid sync or multi-kernel. The multi-kernel approach is cleaner. Actually, here's a different idea: Use ` Cooperative Groups` to do grid-level reduction. But that's complex. Let me just go with 3 kernels and a simple approach: 1. Compute partial stats 2. Reduce to mean and rstd (custom small kernel) 3. Normalize Actually, I just realized: if we pass `mean` and `rstd` as float arrays to the normalize kernel, and these are computed by a custom reduce kernel launched on the same stream, then CUDA stream ordering ensures correctness. No sync issues! So the sequence is: ```cpp partial_stats_kernel<<<grid1, block1, 0, stream>>>(...); reduce_stats_kernel<<<batch_size, block2, 0, stream>>>(...); normalize_kernel<<<grid1, block1, 0, stream>>>(...); ``` This is 3 kernel launches on the same stream. Perfect. But wait, reduce_stats_kernel takes the output of partial_stats_kernel. Since they're on the same stream, the second kernel won't start until the first completes. And normalize won't start until reduce completes. This is exactly what we want. And we don't need `at::sum()` or anything. Pure CUDA. Let me design reduce_stats_kernel: ```cpp extern "C" __global__ void reduce_stats_kernel( const double* __restrict__ partial_sums, const double* __restrict__ partial_sum_sqs, float* __restrict__ mean, float* __restrict__ rstd, int batch_size, int blocks_per_batch, int N, float eps) { int batch_idx = blockIdx.x; if (batch_idx >= batch_size) return; double sum = 0.0; double sum_sq = 0.0; int base = batch_idx * blocks_per_batch; for (int i = threadIdx.x; i < blocks_per_batch; i += blockDim.x) { sum += partial_sums[base + i]; sum_sq += partial_sum_sqs[base + i]; } // Block reduction with warp shuffle #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); } __shared__ double s_sum[32]; __shared__ double s_sum_sq[32]; int lane = threadIdx.x & 31; int warp_id = threadIdx.x >> 5; if (lane == 0) { s_sum[warp_id] = sum; s_sum_sq[warp_id] = sum_sq; } __syncthreads(); if (warp_id == 0) { sum = (lane < (blockDim.x + 31) / 32) ? s_sum[lane] : 0.0; sum_sq = (lane < (blockDim.x + 31) / 32) ? s_sum_sq[lane] : 0.0; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); } if (lane == 0) { double n_d = static_cast<double>(N); double m = sum / n_d; double var = sum_sq / n_d - m * m; double rstd_val = rsqrt(var + static_cast<double>(eps)); mean[batch_idx] = static_cast<float>(m); rstd[batch_idx] = static_cast<float>(rstd_val); } } } ``` For `warp_id == 0`, we need `(blockDim.x + 31) / 32` which is 8 for blockDim.x=256. Let me simplify: for this small kernel, I can just use blockDim.x = 256 and use the exact count. Wait, since `blocks_per_batch` is at most 1024, and we're using blockDim.x = 256, each thread processes up to 4 partial sums. The reduction is fast. Actually, I can make this even simpler. Since blocks_per_batch is typically 64-256, I could use a loop and no inter-thread reduction: each thread just processes a portion, then use `atomicAdd` to a shared location. But atomics on double are not available on all GPU architectures. On compute capability 7.0+, `atomicAdd(double*, double)` is available. H100 is compute capability 9.0, so it's fine. Even simpler: use atomicAdd to accumulate directly in the stats kernel! Instead of writing partial sums and reducing them, each block can `atomicAdd(&global_sum[batch_idx], block_sum)`. Then we don't need a reduce kernel at all! But we'd still need to know when all blocks are done before computing rstd. Unless we launch a separate tiny kernel that reads the atomically accumulated values and computes rstd. Wait, we can do: 1. Stats kernel: each block computes sum and sum_sq, atomically adds to global_sum[batch_idx] and global_sum_sq[batch_idx] 2. Reduce/compute kernel: one thread per batch reads global_sum, computes mean and rstd This is just 2 kernels! And we avoid workspace allocation for partials (except global_sum and global_sum_sq which are tiny: batch_size elements each). But atomicAdd on double for 256 blocks per batch... each block does 2 atomics. 512 atomics per batch, hitting 2 addresses. With 16 batches, that's 16 addresses. The atomics from different batches don't conflict. Within one batch, 256 atomics to the same address. This could have some contention, but H100 has very fast atomics via L2. And it's only 512 atomics total, which is negligible compared to the millions of float operations. Actually, let me think about whether atomics on double are fast enough. On modern GPUs, atomicAdd on global memory for double goes through L2 and is serialized for the same address across a wave/scheduler. For 256 atomicAdds from different SMs to the same address, it might take some time, but the time is on the order of microseconds. Compared to the kernel time of milliseconds, this is tiny. Let's verify: if the main kernel processes 16 * 4M = 64M floats, at 3.35 TB/s, the memory bandwidth alone is 256 MB / 3.35 TB/s ≈ 0.076 ms. The actual kernel with compute might take 0.1-0.2 ms. Adding a few microseconds of atomic contention is 1-2% overhead. That's acceptable. And we save memory by not allocating partial stats arrays. Actually, for batch_size=16, even with 256 blocks_per_batch, the partial stats array is 16*256*16 bytes = 65KB. That's tiny. But with atomics, we only need 16 * 2 * 8 bytes = 256 bytes. Even tinier. And we save the reduce kernel launch overhead. So 2 kernels total: 1. Stats kernel (with atomicAdd) 2. Normalize kernel Let me design this. **Stats kernel with atomics:** ```cpp extern "C" __global__ void compute_mean_rstd_kernel( const float* __restrict__ input, double* __restrict__ global_sum, double* __restrict__ global_sum_sq, float* __restrict__ mean, float* __restrict__ rstd, int batch_size, int N, int blocks_per_batch, float eps) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; int tid = threadIdx.x; int total_threads_per_batch = blocks_per_batch * blockDim.x; int global_tid = block_in_batch * blockDim.x + tid; if (batch_idx >= batch_size) return; const float* batch_input = input + batch_idx * (size_t)N; double sum = 0.0; double sum_sq = 0.0; int num_float4 = N / 4; for (int i = global_tid; i < num_float4; i += total_threads_per_batch) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; sum += val.x; sum_sq += val.x * val.x; sum += val.y; sum_sq += val.y * val.y; sum += val.z; sum_sq += val.z * val.z; sum += val.w; sum_sq += val.w * val.w; } int remainder_start = num_float4 * 4; for (int i = global_tid + remainder_start; i < N; i += total_threads_per_batch) { float val = batch_input[i]; sum += val; sum_sq += val * val; } // Block reduction #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); } __shared__ double s_sum[32]; __shared__ double s_sum_sq[32]; int lane = tid & 31; int warp_id = tid >> 5; if (lane == 0) { s_sum[warp_id] = sum; s_sum_sq[warp_id] = sum_sq; } __syncthreads(); if (warp_id == 0) { sum = (lane < 8) ? s_sum[lane] : 0.0; sum_sq = (lane < 8) ? s_sum_sq[lane] : 0.0; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); } if (lane == 0) { atomicAdd(&global_sum[batch_idx], sum); atomicAdd(&global_sum_sq[batch_idx], sum_sq); } } } ``` Wait, but after all blocks atomically add, how do we know when it's safe to compute mean and rstd? We'd need a grid sync or launch a second kernel. So we still need a second kernel to compute mean and rstd from global_sum and global_sum_sq. But this second kernel is extremely simple: just a few operations per batch. ```cpp extern "C" __global__ void finalize_stats_kernel( const double* __restrict__ global_sum, const double* __restrict__ global_sum_sq, float* __restrict__ mean, float* __restrict__ rstd, int batch_size, int N, float eps) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= batch_size) return; double n_d = static_cast<double>(N); double m = global_sum[idx] / n_d; double var = global_sum_sq[idx] / n_d - m * m; double rstd_val = rsqrt(var + static_cast<double>(eps)); mean[idx] = static_cast<float>(m); rstd[idx] = static_cast<float>(rstd_val); } ``` This kernel is launched with `grid=(1,), block=(batch_size,)` or `grid=((batch_size+255)/256,), block=(256,)`. For batch_size=16, just 1 block of 16 threads. And before this kernel, we need to make sure all the atomics from the stats kernel are visible. But since we launch them on the same stream, the second kernel won't start until the first completes. CUDA stream ordering ensures all writes from kernel 1 are visible to kernel 2. So the sequence is: 1. Zero `global_sum` and `global_sum_sq` (need to memset or initialize to 0) 2. Launch stats kernel with atomics 3. Launch finalize stats kernel 4. Launch normalize kernel Total: 3 kernels. Actually, for step 1, we can allocate and zero the buffers using `cudaMemsetAsync` or just allocate with `torch::zeros`. Let's use `at::zeros`. Now, global_sum and global_sum_sq arrays have size batch_size each. On double, that's 16 * 8 = 128 bytes each. Very small. Wait, the finalize stats kernel needs to read global_sum and global_sum_sq from the same stream, so it's fine. But there's a subtle issue: when we do `atomicAdd` in kernel 1 and then kernel 2 reads the results, since they're on the same stream, it's guaranteed to work. No additional sync needed. Actually, this is great. But I need to zero the buffers first. If I use `torch.zeros`, it's fine. The atomicAdd just accumulates into zero. Let me think about whether 3 kernels is still faster than PyTorch's implementation. PyTorch likely uses cuDNN or a fused kernel. cuDNN's layer norm might use a single kernel or highly optimized multiple passes. To potentially beat PyTorch, we need to be memory-bandwidth-optimal. Our approach: - Kernel 1: reads input once - Kernel 3: reads input, weight, bias once; writes output once - Total read: 2 * input_size + weight_size + bias_size = 2 * 256MB + 1MB + 1MB ≈ 258MB - Total write: output = 256MB - Total memory traffic ≈ 514MB - At 3.35 TB/s: 514 / 3350 = 0.153 ms theoretical If we can get close to this, we'd be very fast. PyTorch might have similar memory traffic but with kernel launch overhead and possibly less optimal vectorization. Wait, actually if PyTorch uses cuDNN's LayerNorm, it might use Tensor Cores or have other optimizations. But for fp32, there's no Tensor Core benefit for LayerNorm. The main thing is memory bandwidth. However, I just realized: our stats kernel with double accumulation might not be bandwidth-bound at all because double FMAs have lower throughput than float loads. Actually, H100 has full-rate fp32 and half-rate fp64. With 4M elements per batch and 16 batches, each doing a few double ops per element, we have: - Loads: 64M floats = 256MB - Double ops: maybe 4 per element = ~1B ops. At 34 TFLOPS fp64, that's ~0.03 ms. So negligible. Actually, each element is loaded once and contributes to two accumulators (sum and sum_sq), with a total of maybe 6-8 double ops per 4 elements. So the compute is tiny. For the normalize kernel: - Loads: input (256MB) + weight (1MB) + bias (1MB) - Stores: output (256MB) - Compute: ~10 fp32 ops per element = ~670M ops. At 67 TFLOPS, that's 0.01 ms. Negligible. - Total memory: ~514MB So we're heavily bandwidth bound. The theoretical minimum is ~0.15ms. In practice, with overhead, maybe 0.2-0.3ms. Is PyTorch faster than 0.2ms for this? Let me guess. For a 256MB tensor, just reading it takes 0.076ms. PyTorch's LayerNorm might take 0.2-0.5ms. So if we can achieve near-memory-bandwidth performance, we'll win. One concern: our 3 kernels approach might have launch overhead. Each kernel launch is ~1-3 microseconds. 3 launches = ~5-10µs, which is negligible compared to 200µs. Another concern: the `atomicAdd` in the stats kernel. For 256 blocks per batch × 16 batches = 4096 blocks, each doing 2 atomicAdds. Total: 8192 atomicAdds. Each atomic requires round-trips to L2/global memory. But H100 has extremely fast L2 atomics. And we only touch 16×2=32 addresses total (with all atomics to the same address being serialized per batch). Still, 256 atomicAdds per address might take a few microseconds. That's fine. Let me think about whether we can avoid atomics and still use just 2 kernels. Alternative: Use a single workspace array of size `batch_size * blocks_per_batch` for partial sums. Launch the stats kernel writing to workspace. Then compute mean and rstd directly on CPU or GPU. But we still need to launch a second kernel or do a PyTorch reduction, then a third kernel for normalize. That's still 3 kernels. Unless we do the reduction on GPU in the same kernel as normalize, which is impossible because normalize needs rstd. Or we can use the persistent block pattern: have some blocks do reduction while others normalize. But that's complex. I think 3 kernels is fine. Let me code it up. Actually, another idea: what if we fuse the finalize stats kernel into the normalize kernel? We can't, because normalize needs rstd across all elements. Wait, could we have normalize kernel first compute the sum of squared deviations in a local reduction? No, because we need the global mean first, which requires global reduction. Actually, there's an algorithm called "parallel sum then normalize" that does it in 2 kernels: 1. Compute mean: reduction kernel 2. Compute variance + normalize in single kernel: each block computes local sum of (x-mean)^2, reduces to block-level, all blocks do atomicAdd to global_sum_sq, then normalize. But the problem is: all blocks need the global variance before normalizing. Blocks within a batch can't normalize until all blocks have contributed to the global variance. With persistent blocks and cooperative grid sync, you could do it in 1 kernel. Without that, you need multiple kernels. Actually, I wonder if we can use `cudaDeviceSynchronize()` between grid launches in a single kernel? No, that's not possible. So 3 kernels is the minimum for our approach. Let me proceed. Wait, actually there's a simpler 2-kernel approach: 1. Kernel: compute mean. Use single-pass tree reduction with multiple blocks per batch writing to global memory. For this, each block computes partial sum, and we use atomicAdd to a single mean accumulator per batch. Then have a second phase in the SAME kernel where threads check if they're the "finalization" thread and compute rstd... but this isn't safe without inter-block sync. Or with a 2-kernel approach: 1. Kernel computes mean (reduction) 2. Kernel computes (x-mean)^2 sum via atomicAdd to global variance accumulator, AND normalizes each element. For kernel 2, the issue is that the first blocks to finish will start normalizing before all blocks have contributed to the variance. This gives incorrect results. Unless we use 2 separate launches: - Kernel 2a: compute variance (sum of (x-mean)^2) - Kernel 2b: normalize That's back to 3 kernels. OK, 3 kernels it is. Let me finalize the code structure. One more consideration: thread block size. I'll use 256 threads for both main kernels and 16-256 threads for the tiny finalize kernel. For compute_mean_rstd_kernel, I need to zero the accumulators before the atomics. I'll create buffers initialized to 0. For the wrapper C++ function in `cuda_sources`: ```cpp torch::Tensor custom_layer_norm_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, float eps) { int batch_size = input.size(0); int N = input.numel() / batch_size; // Allocate output and intermediate buffers auto output = torch::empty_like(input); auto mean = torch::empty({batch_size}, torch::dtype(torch::kFloat32).device(input.device())); auto rstd = torch::empty({batch_size}, torch::dtype(torch::kFloat32).device(input.device())); auto global_sum = torch::zeros({batch_size}, torch::dtype(torch::kFloat64).device(input.device())); auto global_sum_sq = torch::zeros({batch_size}, torch::dtype(torch::kFloat64).device(input.device())); int threads_per_block = 256; int elements_per_thread = 64; int blocks_per_batch = std::max(1, (N + threads_per_block * elements_per_thread - 1) / (threads_per_block * elements_per_thread)); blocks_per_batch = std::min(blocks_per_batch, 1024); dim3 grid(batch_size * blocks_per_batch); dim3 block(threads_per_block); auto stream = at::cuda::getCurrentCUDAStream(input.device().index()); compute_stats_kernel<<<grid, block, 0, stream.stream()>>>( input.data_ptr<float>(), global_sum.data_ptr<double>(), global_sum_sq.data_ptr<double>(), batch_size, N, blocks_per_batch ); finalize_stats_kernel<<<batch_size, 1, 0, stream.stream()>>>( global_sum.data_ptr<double>(), global_sum_sq.data_ptr<double>(), mean.data_ptr<float>(), rstd.data_ptr<float>(), batch_size, N, eps ); normalize_kernel<<<grid, block, 0, stream.stream()>>>( input.data_ptr<float>(), weight.defined() ? weight.data_ptr<float>() : nullptr, bias.defined() ? bias.data_ptr<float>() : nullptr, output.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(), batch_size, N, blocks_per_batch ); return output; } ``` Wait, there's a bug in `finalize_stats_kernel`: I should probably launch it with more than 1 thread per block for better performance, but since batch_size=16, it barely matters. `<<< (batch_size + 255) / 256, 256 >>>` would be more standard. But 1 block with 256 threads where only 16 are active is fine. Actually, let me use `dim3((batch_size + 255) / 256, 1, 1)` and `dim3(256, 1, 1)` to be more standard. For the finalize kernel code: ```cpp extern "C" __global__ void finalize_stats_kernel(...) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= batch_size) return; ... } ``` One more thing: when we pass `weight` and `bias` as optional tensors, we need to handle the case where they're undefined. In PyTorch C++, `torch::Tensor` can be undefined. We use `weight.defined()` to check. But wait, what if `weight` and `bias` are None in Python? When passed to the C++ wrapper, they become undefined tensors. So `weight.defined()` returns false. Good. Let me also handle the case where `input` might not be contiguous. We should probably call `.contiguous()` on it. But this adds overhead. The reference code uses `torch.rand` which creates contiguous tensors by default, and for typical use, inputs are contiguous. To be safe, we can check and make contiguous if needed. ```cpp auto input_contig = input.contiguous(); ``` Similarly for weight and bias. Now, writing the complete Python file... Actually, let me also think about whether to use `__ldg` for loading. For read-only memory, `__ldg` can use a read-only cache. Let me use it for input in both kernels: ```cpp float4 val = __ldg(reinterpret_cast<const float4*>(batch_input) + i); ``` Actually, `__ldg` doesn't work with `float4` directly in all CUDA versions/ architectures. But on modern CUDA, it should. To be safe, I'll stick with regular loads, which are already very efficient for linear access. Let me write one more optimization: in the compute stats kernel, after loading the float4, we can process all 4 elements before moving to the next float4. This is what my current code does. For the normalize kernel, I'd use similar structure. Actually, let me also consider whether we need `size_t` for indexing. With N=4,194,304 and batch_size=16, total elements = 67M, which fits in int. But for safety, using size_t for pointer arithmetic is good. OK, let me write the kernel code now. I'll put everything in `cuda_sources` (the host wrapper + device kernels). Wait, `cuda_sources` needs `extern "C"` for the CUDA kernels if we want to call them from `cpp_sources`, but since we're putting the wrapper in `cuda_sources` too, they can just be regular global functions. But to avoid C++ name mangling issues, I'll use `extern "C"` for the kernels. Actually, the wrapper function in `cuda_sources` is compiled as C++ with nvcc, so it can see the `__global__` functions directly. No need for `extern "C"` if they're in the same file. But `extern "C"` is safer and good practice. Let me draft out the complete code: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline cuda_source = r""" #include <torch/extension.h> #include <cuda_runtime.h> #include <cmath> __global__ void compute_stats_kernel( const float* __restrict__ input, double* __restrict__ global_sum, double* __restrict__ global_sum_sq, int batch_size, int N, int blocks_per_batch) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; int tid = threadIdx.x; int total_threads_per_batch = blocks_per_batch * blockDim.x; int global_tid = block_in_batch * blockDim.x + tid; if (batch_idx >= batch_size) return; const float* batch_input = input + static_cast<size_t>(batch_idx) * N; double sum = 0.0; double sum_sq = 0.0; // Coalesced float4 loads int num_float4 = N / 4; for (int i = global_tid; i < num_float4; i += total_threads_per_batch) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; sum += val.x; sum_sq += val.x * val.x; sum += val.y; sum_sq += val.y * val.y; sum += val.z; sum_sq += val.z * val.z; sum += val.w; sum_sq += val.w * val.w; } // Tail int remainder_start = num_float4 * 4; for (int i = global_tid + remainder_start; i < N; i += total_threads_per_batch) { float val = batch_input[i]; sum += val; sum_sq += val * val; } // Warp-level reduction #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); } __shared__ double s_sum[32]; __shared__ double s_sum_sq[32]; int lane = tid & 31; int warp_id = tid >> 5; if (lane == 0) { s_sum[warp_id] = sum; s_sum_sq[warp_id] = sum_sq; } __syncthreads(); // Final reduction by warp 0 if (warp_id == 0) { sum = (lane < ((blockDim.x + 31) >> 5)) ? s_sum[lane] : 0.0; sum_sq = (lane < ((blockDim.x + 31) >> 5)) ? s_sum_sq[lane] : 0.0; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); } if (lane == 0) { atomicAdd(&global_sum[batch_idx], sum); atomicAdd(&global_sum_sq[batch_idx], sum_sq); } } } __global__ void finalize_stats_kernel( const double* __restrict__ global_sum, const double* __restrict__ global_sum_sq, float* __restrict__ mean, float* __restrict__ rstd, int batch_size, int N, float eps) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= batch_size) return; double n = static_cast<double>(N); double m = global_sum[idx] / n; double var = global_sum_sq[idx] / n - m * m; double rstd_val = rsqrt(var + static_cast<double>(eps)); mean[idx] = static_cast<float>(m); rstd[idx] = static_cast<float>(rstd_val); } __global__ void normalize_kernel( const float* __restrict__ input, const float* __restrict__ weight, const float* __restrict__ bias, float* __restrict__ output, const float* __restrict__ mean, const float* __restrict__ rstd, int batch_size, int N, int blocks_per_batch) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; int tid = threadIdx.x; int total_threads_per_batch = blocks_per_batch * blockDim.x; int global_tid = block_in_batch * blockDim.x + tid; if (batch_idx >= batch_size) return; float m = mean[batch_idx]; float rs = rstd[batch_idx]; const float* batch_input = input + static_cast<size_t>(batch_idx) * N; float* batch_output = output + static_cast<size_t>(batch_idx) * N; bool has_weight = (weight != nullptr); bool has_bias = (bias != nullptr); int num_float4 = N / 4; for (int i = global_tid; i < num_float4; i += total_threads_per_batch) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; float4 w = has_weight ? reinterpret_cast<const float4*>(weight)[i] : make_float4(1.0f, 1.0f, 1.0f, 1.0f); float4 b = has_bias ? reinterpret_cast<const float4*>(bias)[i] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); val.x = (val.x - m) * rs * w.x + b.x; val.y = (val.y - m) * rs * w.y + b.y; val.z = (val.z - m) * rs * w.z + b.z; val.w = (val.w - m) * rs * w.w + b.w; reinterpret_cast<float4*>(batch_output)[i] = val; } int remainder_start = num_float4 * 4; for (int i = global_tid + remainder_start; i < N; i += total_threads_per_batch) { float w = has_weight ? weight[i] : 1.0f; float b_val = has_bias ? bias[i] : 0.0f; batch_output[i] = (batch_input[i] - m) * rs * w + b_val; } } torch::Tensor custom_layer_norm_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, float eps) { auto input_contig = input.contiguous(); int batch_size = input_contig.size(0); int N = input_contig.numel() / batch_size; auto output = torch::empty_like(input_contig); auto mean = torch::empty({batch_size}, torch::TensorOptions().dtype(torch::kFloat32).device(input_contig.device())); auto rstd = torch::empty({batch_size}, torch::TensorOptions().dtype(torch::kFloat32).device(input_contig.device())); auto global_sum = torch::zeros({batch_size}, torch::TensorOptions().dtype(torch::kFloat64).device(input_contig.device())); auto global_sum_sq = torch::zeros({batch_size}, torch::TensorOptions().dtype(torch::kFloat64).device(input_contig.device())); int threads_per_block = 256; int elements_per_thread = 64; int blocks_per_batch = std::max(1, (N + threads_per_block * elements_per_thread - 1) / (threads_per_block * elements_per_thread)); blocks_per_batch = std::min(blocks_per_batch, 1024); // Pad batch_size * blocks_per_batch to be enough for the grid dim3 grid(batch_size * blocks_per_batch); dim3 block(threads_per_block); cudaStream_t stream = at::cuda::getCurrentCUDAStream(input_contig.device().index()); compute_stats_kernel<<<grid, block, 0, stream>>>( input_contig.data_ptr<float>(), global_sum.data_ptr<double>(), global_sum_sq.data_ptr<double>(), batch_size, N, blocks_per_batch ); finalize_stats_kernel<<<(batch_size + 255) / 256, 256, 0, stream>>>( global_sum.data_ptr<double>(), global_sum_sq.data_ptr<double>(), mean.data_ptr<float>(), rstd.data_ptr<float>(), batch_size, N, eps ); normalize_kernel<<<grid, block, 0, stream>>>( input_contig.data_ptr<float>(), weight.defined() ? weight.data_ptr<float>() : nullptr, bias.defined() ? bias.data_ptr<float>() : nullptr, output.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(), batch_size, N, blocks_per_batch ); return output; } """ custom_ln = load_inline( name="custom_ln", cpp_sources="", cuda_sources=cuda_source, functions=["custom_layer_norm_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False, ) class ModelNew(nn.Module): def __init__(self, normalized_shape): super().__init__() self.normalized_shape = tuple(normalized_shape) if not isinstance(normalized_shape, tuple) else normalized_shape self.eps = 1e-5 self.weight = nn.Parameter(torch.ones(self.normalized_shape)) self.bias = nn.Parameter(torch.zeros(self.normalized_shape)) def forward(self, x): return custom_ln.custom_layer_norm_forward( x, self.weight, self.bias, self.eps ) ``` Wait, there's a problem: `blocks_per_batch = (N + threads_per_block * elements_per_thread - 1) / (threads_per_block * elements_per_thread)` can be 0 when N < threads_per_block * elements_per_thread. But I have `std::max(1, ...)`, so it handles that. For N = 4,194,304: - threads_per_block * elements_per_thread = 256 * 64 = 16,384 - blocks_per_batch = (4,194,304 + 16,383) / 16,384 = 4,210,687 / 16,384 = 257.000... wait! 4,194,304 / 16,384 = 256 exactly. (4,194,304 + 16,383) / 16,384 = 4,210,687 / 16,384 = 257.00006... which floors to 257 in integer division. But wait, do we want exactly 256 or is 257 fine? 257 is prime, so grid-stride loops with 257 blocks and 256 threads would give total_threads_per_batch = 65,792. Then N/4 = 1,048,576 float4s. 1,048,576 / 65,792 ≈ 15.938. So some threads would iterate 16 times and others 15 times. That's fine, but the workload might not divide evenly. Actually, the exact formula: (X + Y - 1) / Y gives ceil(X / Y). For power-of-2 X and Y, if X is divisible by Y, then (X + Y - 1) / Y = X/Y + (Y-1)/Y. In integer division: (X + Y - 1) / Y = X/Y because X is exactly divisible by Y and Y-1 < Y. Wait, X=4,194,304, Y=16,384. X/Y = 256. (X + Y - 1) / Y = (4,194,304 + 16,383) / 16,384 = 4,210,687 / 16,384. 4,210,687 = 256 * 16,384 + 16,383. So 4,210,687 / 16,384 = 256 with remainder 16,383. Integer division gives 256. Actually, in C and Python 3 integer division: 4194304 // 16384 = 256 (4194304 + 16383) // 16384 = 4210687 // 16384 4210687 / 16384: 16384 * 256 = 4194304 4210687 - 4194304 = 16383 So 4210687 = 16384 * 256 + 16383 Integer division gives 256. So we get exactly 256! Perfect. What if X = 4,194,300 (not divisible by 4)? Then we'd get blocks_per_batch = ceil(4,194,300 / 16,384) = ceil(256.00024) = 257. That's fine. Now, for blocks_per_batch = 257 and total_threads_per_batch = 257 * 256 = 65,792: N/4 = 1,048,575 (since N=4,194,300) 1,048,575 / 65,792 ≈ 15.935 So threads 0 to (1,048,575 % 65,792 - 1) would process 16 float4s, and the rest 15. That's fine. But the indexing in the normalize kernel if we use `weight[i]` where `i` is a float4 index: for `num_float4 = N / 4` (truncating division), the float4 elements in the tail would be left over. Wait, the tail loop handles `remainder_start = num_float4 * 4` which covers all leftover elements. For N not divisible by 4, the remainder is 1-3 elements and they are handled by the tail loop. That's correct. But if N is less than 4, then `num_float4 = 0` and the tail loop handles all elements. Good. There's one thing I need to double-check in the indexing for weight. In the normalize kernel, for the float4 loop, `i` ranges from 0 to `num_float4 - 1`. We load `weight[i]` as a float4. This assumes `weight` is at least `num_float4 * 4` elements, which it is since weight has shape [N] and `num_float4 * 4 <= N`. That's correct. For the tail loop, `i` ranges from `num_float4 * 4` to `N - 1`. We load `weight[i]`. That's also correct since `i < N`. Great. One more potential issue: what if `batch_size * blocks_per_batch` exceeds the maximum grid dimension? CUDA allows grid.x up to 2^31 - 1. batch_size is 16 and blocks_per_batch is at most 1024, so grid.x = 16384. That's tiny. Another issue: when input is not contiguous or has weird strides, `input.contiguous()` fixes it but allocates a new tensor. For the reference input which is `torch.rand()`, it's already contiguous. So we're fine. Now I should also think about numerical issues once more. If the data has mean 0.5 and variance 1/12, then: - `var = sum_sq / N - m * m` For fp64 accumulation, this is accurate to ~1e-16. For fp32 output with tolerance 1e-4, we're definitely fine. But what about cases where the variance is very small or there's catastrophic cancellation? For random uniform data, it's fine. For adversarial inputs, we might lose precision. But the reference is PyTorch's layer norm, which also uses some floating point accumulation. With fp64, our accumulation should be MORE accurate than PyTorch's, which likely uses fp32 Welford or similar. Wait, but being more accurate might actually make us DIFFERENT from the reference! The tolerance is 1e-4 though, so even if we're more accurate, the difference should be within tolerance. However, there is a subtlety: PyTorch computes `rstd = 1.0 / sqrt(var + eps)`. Different algorithms for `var` could give slightly different results. But with fp64, our var should be closer to the "true" mathematical variance than PyTorch's fp32 computation. The question is whether the difference exceeds 1e-4 absolute. For typical values: if x is in [0,1], then normalized x is roughly z = (x - 0.5) / 0.289. That's around [-1.7, 1.7]. A difference in var of 1e-6 would change rstd by about 1e-6, changing z by about 1e-6, which is well within 1e-4. So I think we're safe. But let me also consider using `m = sum / N` and then `var = (sum_sq - m * sum) / N` instead of `sum_sq / N - m * m`. These are algebraically equivalent but differ numerically. The second form `sum_sq / N - m * m` is standard but can suffer catastrophic cancellation if var << m^2. For our case, var ≈ 0.083 and m^2 ≈ 0.25, so var is not much smaller than m^2. The relative error is acceptable. The alternative `var = (sum_sq - m * sum) / N = (sum_sq - sum^2 / N) / N = (N * sum_sq - sum^2) / N^2`. This is equivalent to the "two-pass corrected" formula. Let me just use `var = sum_sq / N - m * m` with fp64. It should be accurate enough. Wait, one more thing. `atomicAdd` on double requires compute capability 6.0+. H100 is 9.0, so it's fine. But what if the code is tested on an older GPU? Since the environment is H100, we're good. However, `atomicAdd` for double uses the function: ```cpp double atomicAdd(double* address, double val); ``` This is defined in `cuda_runtime.h` for sm_60+. So we don't need to define it ourselves. Let me compile and test this. Actually, since I can only call compile_kernel and other tools, let me write the complete Python code and compile it. Wait, I need to be careful with `load_inline`. For the C++ wrapper, I should make sure the signature is right. If I write it in `cuda_sources`, it will be compiled as CUDA C++ and the wrapper function should be callable from Python. Also, `torch::Tensor` in `cuda_sources`: need to make sure we have the right includes. `torch/extension.h` should provide everything. One more issue: `at::cuda::getCurrentCUDAStream(input.device().index())` returns a `c10::cuda::CUDAStream`, and we can get the raw `cudaStream_t` with `.stream()`. But I need to include the right headers. Actually, with `torch/extension.h`, the ATen CUDA headers should be included. Let me use the simpler approach: Actually, simpler than managing streams explicitly: we can just not pass a stream and let CUDA use the default stream. PyTorch operations and our kernel launches will all use the default stream, so they're implicitly synchronized in order. For `load_inline` with auto-generated wrappers, does it handle streams automatically? The auto-generated pybind11 wrapper calls the C++ function, which can launch kernels on any stream. But if we don't specify a stream, kernels use stream 0 (default). If PyTorch is using a different stream for its operations (e.g., in DataParallel or with `torch.cuda.stream()`), there could be issues. But `torch.cuda.synchronize()` would typically handle this at boundaries. To be robust, using `at::cuda::getCurrentCUDAStream()` is better. But I need to check if we can call it from our code. Actually, `torch/extension.h` usually includes `ATen/cuda/CUDAContext.h` indirectly. Alternatively, I use `cudaStream_t stream = c10::cuda::getCurrentCUDAStream(input.device().index());` Hmm, let me just not worry about streams and use the default stream for simplicity. In practice, PyTorch's eager mode typically uses the default stream anyway. So my wrapper will launch kernels without explicit streams: ```cpp compute_stats_kernel<<<grid, block>>>(...); finalize_stats_kernel<<<(batch_size + 255) / 256, 256>>>(...); normalize_kernel<<<grid, block>>>(...); ``` Kernel launches on the same (default) stream are ordered sequentially. OK let me write the code and see if it compiles. I'll create the source code and call compile_kernel. Actually, one more thought: `make_float4` is defined in `vector_types.h` which is included by `cuda_runtime.h`. But `make_float4(1.0f, 1.0f, 1.0f, 1.0f)` should work. However, sometimes `make_float4` doesn't work in device code with `load_inline` due to some compiler quirks. To be safe, I can use: ```cpp float4 w = make_float4(1.0f, 1.0f, 1.0f, 1.0f); ``` If that fails, I can use: ```cpp float4 w; w.x = 1.0f; w.y = 1.0f; w.z = 1.0f; w.w = 1.0f; ``` Let me use a helper function or inline initialization. Actually, in C++, we can also do: ```cpp float4 w{1.0f, 1.0f, 1.0f, 1.0f}; ``` For `make_float4` in `__global__` code, it should be available. Let me just use it and see if there's a compilation issue. Wait, another thing: load_inline creates a temporary directory and compiles everything. The compilation might take a few seconds, but that's fine since it's done at module load time. Let me write the full code and test compile.

The compilation failed because `load_inline` needs the function declaration to be visible in the same translation unit. When using `cpp_sources=""`, the CUDA code and the auto-generated binding code are in separate files. I need to add a header declaration. Actually, I should put the function declaration in cpp_sources and the implementation in cuda_sources. Or I can put everything in a single approach. Let me fix this by adding the function declaration in the cpp_sources. Actually the standard approach with `load_inline` is to put function declarations in `cpp_sources` with `torch::Tensor myfunc(...)` and the definitions in `cuda_sources`. Let me do that.

Good! The kernel compiled. Now let me check correctness.

The correctness check passed. Now let me run the static check and submit the kernel. But first, let me think if there are any optimizations I can make. Actually, let me check the static_analysis and then just submit. The numerical precision looks good. But wait - I should think about whether my kernel is actually fast. Let me analyze: 1. For batch_size=16, N=4194304: 2. Stats kernel: 16 batches * 256 blocks = 4096 blocks total 3. Normalize kernel: same 4096 blocks The stats kernel reads all elements once (of the input), and the normalize kernel reads all three inputs (input, weight, bias) and writes the output. So theoretically we need to read 5 * N bytes and write 1 * N bytes = ~420MB. With 3.35 TB/s bandwidth, that would take about 125 microseconds. But we need two passes over the data (stats + normalize), so maybe ~250us. This is likely memory bandwidth bound. Actually, a fused kernel that computes stats on-the-fly and writes in one pass could be faster. But since we need all elements for the mean/variance, we can't truly fuse it in one pass. However, we can optimize by: 1. Using a single kernel that does a persistent approach: load a block's data, compute local partial stats, synchronize, compute final mean/rstd, then normalize and write back 2. This would reduce global memory bandwidth by almost half Let me try to implement a fused two-pass approach where we: 1. First pass: compute stats per block (persistent state in shared memory), then within each batch, reduce to get mean/rstd, store in shared memory 2. Second phase: normalized and write back Actually, the most efficient approach for layer norm on large feature dimensions is to use cooperative groups and persistent kernels. Let me try a different approach. Let me think about this differently. For each batch of 4M elements: - We need to read all elements to compute sum and sum_sq - Then we need to read all elements again to normalize them With float4, each thread processes 4 elements per iteration. For the stats phase, each thread reads 4 elements, accumulates, then we reduce across the block and atomically add to global. For the normalize phase, each thread reads 4 elements, reads weight and bias, writes normalized output. Total memory traffic per element: - Stats: read input (4 bytes) - Normalize: read input (4), read weight (4), read bias (4), write output (4) = 16 bytes - Total: 20 bytes per element Wait, actually the stats also needs to read all elements. And the normalize needs to read again. So it's: - Stats: ~16 bytes (both input and weight/bias are read once per element, but weight/bias for stats is not needed, only input) - Normalize: 16 bytes as above - Total: 20 bytes Hmm, but we read input twice (for stats and normalize). So 8 + 16 = 24 bytes per element. For 67M elements that's ~1.6GB of data movement. With 3.35 TB/s bandwidth and 80% efficiency, that's ~1000 GB/s effective. So 1.6 GB / 1000 GB/s = 1.6 ms. Actually I think our kernel will be much slower because we're launching many blocks and there's overhead. Let me consider optimizing further with a fused approach. Fused kernel approach: - For each batch, we can have blocks that: 1. Each block reads its portion of data and computes partial sum and sum_sq 2. Use global atomics or a two-phase reduction to get global mean and rstd 3. Synchronize all blocks (using CUDA cooperative groups or atomic counter) 4. Once mean/rstd is known, each block normalizes its portion and writes output This requires careful synchronization. Using cooperative groups could be complex. Let me try a simpler approach: Actually, for H100 with 132 SMs, and we have 16 batches, the ideal case is to have blocks that process entire batches entirely on one SM, avoiding cross-SM communication. But with 4M elements per batch, even at 1 element per thread, we'd need 1024 threads * 4 elements/thread = 4096 per block, so ~1024 blocks needed per batch. Wait, let me reconsider. Maybe I should just use a simpler two-pass approach but make sure to: 1. Coalesce memory accesses properly (float4) 2. Use appropriate block/thread counts 3. Minimize shared memory usage to maximize occupancy My current kernel is probably fine for correctness but might not be super fast. Let me just submit it and see what happens. The key insight is that layer norm on such a large dimension is fundamentally memory bandwidth bound, and the PyTorch reference might also be bandwidth bound. Actually, let me reconsider. PyTorch's native LayerNorm might be using cuDNN or an optimized kernel from ATen. It could be very well optimized already, potentially using a single fused kernel. To beat it, I need to be smarter. Let me try a fully fused kernel where: 1. We allocate intermediate global memory for mean and rstd 2. We launch a grid with exactly batch_size blocks (one per batch) 3. Each block first computes stats cooperatively over all N elements 4. Then normalizes and writes output For N = 4,194,304 elements and batch_size = 16: - Each block handles 4M elements - 256 threads per block means ~16,384 iterations per thread - This is a lot of iterations Instead, maybe use a two-level approach: 1. Have enough blocks to keep all SMs busy for the stats pass 2. Store partial sums in global memory 3. Reduce partial sums to get mean/rstd 4. Normalize in a second pass A better fused approach: - One kernel for stats: launch batch_size * num_blocks_per_batch blocks - Each block computes sum and sum_sq for its assigned elements - Atomic add to global mean and rstd arrays - Second kernel for normalize This is basically what I have. But PyTorch likely does something similar or uses cuDNN which might be faster due to better scheduling. Let me think of another optimization: using Tensor Cores or specialized instructions... no, layer norm is element-wise, no matrix multiplication. Actually, one key optimization could be to compute mean and variance using Welford-style online algorithm which is more numerically stable and can be done in one pass, but for large N it won't make a huge bandwidth difference. Another idea: interleave reading of input, weight, and bias to maximize memory coalescing and cache utilization. In my normalize kernel, reading weight and bias might be perfectly coalesced since all threads in a warp access the same indices. Actually, let me reconsider if the kernel will be faster than PyTorch. PyTorch's LayerNorm for large dimensions might be using: 1. A highly optimized cuDNN/cuBLASLt implementation 2. Or a custom CUDA kernel that's already very efficient For me to beat it, I need to either: 1. Fuse operations that PyTorch doesn't (e.g., avoid intermediate allocations) 2. Use better memory access patterns 3. Reduce kernel launch overhead My kernel uses float4 vectorized loads which is good. Let me try to optimize further by using a single kernel with two passes if possible, but the cross-block synchronization makes that difficult. Actually, maybe I should use aن保cuBLAS-style approach where I have exactly one block per batch. For N = 4M, a block of 1024 threads would need 4096 float4 loads each. That might be too much work per block and would underutilize the GPU (only 16 blocks for 132 SMs). Let me try another approach: use persistent kernels. Each SM stays resident and processes multiple batches. This avoids kernel launch overhead and can better utilize the GPU. For persistent kernels, I can have a grid-stride loop where each block processes multiple batches. But for stats computation, we need to be careful. Another approach: instead of using float4 loads, maybe use __ldg for read-only data to leverage the texture cache. Let me try one more optimization: prefetch data. But that might not help much. Actually, I think the simplest and most likely to be fast approach is what I already have, but let me make sure to tune the parameters: - More blocks per batch for stats could help with parallelism - 1024 threads per block might provide better throughput than 256 Let me try with 512 or 1024 threads and more blocks per batch. Actually, you know what, let me just submit this kernel since it passed correctness. I'll likely need to iterate based on performance feedback. But wait, the performance feedback is only available after submit. Let me just submit and see. Actually no, I should try to optimize further before spending my submission. Let me think about this more carefully. For H100: - 132 SMs - 3.35 TB/s bandwidth - We have 16 batches of 4.19M elements each = 67.1M elements total - Element size: 4 bytes (float32) - Stats pass reads 67.1M * 4 = 268.4 MB - Normalize pass reads 67.1M * 12 (input + weight + bias) and writes 67.1M * 4 = 1.07 GB - Total: ~1.34 GB of memory traffic At 3.35 TB/s, this would take ~400 µs. If we can achieve 80% bandwidth efficiency, that's ~500 µs. PyTorch's implementation is likely doing something similar, but maybe with less overhead. Actually wait, I realize I should check if my weight and bias shapes are correct. In the reference, `nn.LayerNorm(normalized_shape=(64, 256, 256))` creates weight and bias of shape (64, 256, 256). In my kernel, I pass `self.weight` and `self.bias` which should have the same shape. The normalize kernel reads them as flat arrays of size N, which is correct because N = features * dim1 * dim2 = 64 * 256 * 256 = 4,194,304. One more thought: my atomicAdd on double-precision values might be a bottleneck. Instead, I could: 1. Have a two-phase reduction where blocks write partial sums to a temporary buffer 2. A second kernel reduces the partial sums This avoids atomics entirely in the first pass and allows more parallel accumulation. Let me try this: 1. Stats kernel: batch_size blocks, each with enough threads to cover N elements - Each thread computes partial sum and sum_sq - Write to a global scratch buffer 2. Reduce kernel: reduce the scratch buffer per batch 3. Normalize kernel: same as before Actually, for a cleaner approach, let me use a fixed number of blocks per batch (e.g., 128 blocks per batch), each with 256 threads. Each block computes partial sums and writes to a global buffer. Then a second kernel reduces the partial sums per batch. Then a third kernel normalizes. This is basically a 3-pass approach but might be faster because: 1. First pass has no atomics (just writes to global memory) 2. Second pass is a small reduction (16 batches * 128 partial sums) 3. Third pass is the normalize Total memory traffic: input (read 2x) + partial_sums (write+read) + weight + bias + output = ~1.4 GB, similar to before. But avoiding atomics might help. Actually, atomics on fp64 are quite slow on some architectures. For Ampere and Hopper, atomics on shared/global memory for fp64 have decent throughput but still serialize. Let me just keep my current kernel and submit. But first, let me consider if there are any issues with how I handle weight and bias. The reference model uses nn.LayerNorm which has learnable weight and bias parameters. My ModelNew creates these parameters but I need to make sure they have the same names and shapes as the reference. In the reference evaluation, they'll instantiate `ModelNew` with the same `normalized_shape` and pass the same inputs. The state dict of `ModelNew` won't be loaded from the reference, so whether parameter names are the same doesn't matter - only the forward output matters. OK let me submit. Actually wait, let me do a static check first.

All checks pass. Let me now optimize the kernel further to get better speedup: 1. Reduce global memory traffic by computing mean and rstd in shared memory instead of using atomics and a separate kernel 2. Use single-precision accumulation with compensation (Kahan summation or pairwise summation) to match numerically what PyTorch does 3. Fused kernel to compute stats and normalize 4. Vectorized reads/writes (float4) for weight and bias Actually, let me think about PyTorch's implementation: PyTorch uses Welford's algorithm for numerical stability. My current implementation computes mean and variance as E[x] and E[x^2] - E[x]^2, which can have numerical issues but matches closely enough for fp32. Wait, looking at the correctness output: max_rel_err=8.26e+05 might seem high but that's for very small values which is common in layer norm. Let me think about potential issues with my kernel: 1. The float4 alignment: weight and bias parameters are stored contiguously but when they're shaped [features, dim1, dim2], do they have correct alignment for float4 access? In PyTorch, contiguous tensors of shape [features, dim1, dim2] with dim1=256, dim2=256: features * dim1 * dim2 = 64 * 256 * 256 = 4,194,304 elements. A float4 load requires 16-byte (4 float) alignment. The tensor is allocated at a properly aligned address. Since dim2 = 256 is divisible by 4, each row of the last dimension is aligned. But wait - weight and bias have shape [features, dim1, dim2] so they have the same alignment as input. Actually, the input is [batch_size, features, dim1, dim2] = [16, 64, 256, 256]. When I reinterpret this as [batch_size, N] = [16, 4194304], the float4 access alignment is fine since the whole tensor starts at an aligned address and N is divisible by 4. Same for weight and bias of shape [64, 256, 256] - they have 4,194,304 elements aligned. However, let me reconsider: when I compute stats per batch, I read all N = 4,194,304 elements. But weight and bias are shared across the batch dimension. So when I access weight[i] in the normalize kernel, I'm accessing the same elements for all batch items. That's fine. But wait - in the normalize kernel, I reinterpret weight as float4 pointer. That should work fine. Let me consider one more optimization: use warp-level cooperative groups to synchronize warps within a block for reduction instead of shared memory -> this might be slightly faster. Actually, I think there's a better approach to make it faster: 1. Use two-pass approach but with better memory coalescing 2. Avoid global atomics by having each block write its partial sum to global memory, then have a separate small kernel reduce the partial sums Actually the current approach with atomics to double arrays should be fairly efficient for just 16 batches. 3. Use single kernel with persistent blocks or grid stride Actually, I think the current kernel should be reasonably efficient. Let me think about what's potentially slow: - The compute_stats_kernel with atomics: if all blocks for a batch finish at the same time, they serialize on the atomic adds - Three kernel launches instead of one Let me try to improve by: 1. Removing the separate finalize_stats_kernel - this requires launching another small kernel 2. Doing block reduction without atomics, then a second kernel to reduce the per-block partial sums Actually for 16 batches even with hundreds of blocks, atomics to different memory locations (one per batch) should be okay. The issue might be having too many blocks contending for the same batch. Let me try a different approach: use warp-shuffle reduction to compute stats per warp, store warp results in shared memory, then have warp 0 compute final stats and write mean/rstd directly. This avoids global atomics and the extra kernel. But this only works if one block can process all elements of a batch. With 4M elements per batch and 256 threads, each thread would need to process 16K elements, which would require many loop iterations in a single block, but could actually be efficient! Let me try this simpler approach: one block per batch, single kernel that: 1. Uses grid-stride loop to compute sum and sum_sq with float4 loads 2. Does warp + block-level reduction 3. Computes mean and rstd 4. Stores in shared memory 5. Uses grid-stride loop again to normalize with float4 loads This is a fusion of compute_stats + finalize + normalize into one kernel, but each batch gets only one block which limits parallelism if the blocks are small. For 4M elements with 256 threads: - Each thread processes ~16K elements - First pass: sum and sum_sq - Block reduction - Compute mean, rstd - Second pass: normalize The problem is 256 threads might not saturate memory bandwidth for a single batch. With 16 batches, we get 16 * 256 = 4096 threads total, which is fine but each block has to process 4M elements in two passes. Actually, for H100 with 3.35 TB/s bandwidth, reading 16 * 4M * 4 bytes = 256 MB for input, plus 2 * 256 MB for normalize read+write = total ~512 MB. At 3.35 TB/s that's 0.15 ms per pass. With 3 kernel launches we'd be around 0.5ms or so. Fusing into one kernel might reduce overhead. But wait, this model is batch_size = 16 features = 64 dim1 = 256 dim2 = 256 **Total elements per batch** = 64 * 256 * 256 = 4,194,304 ≈ **4.2M** **Float memory reads** (per batch, for stats): - **Read input**: 4.2M floats × 4 bytes = **16.7 MB** **Float memory reads+writes** (for normalize): - **Read input**: 16.7 MB - **Read weight**: 16.7 MB - **Read bias**: 16.7 MB - **Write output**: 16.7 MB - Total: **66.8 MB** For 16 batches: - Stats total read: 16 × 16.7 MB = **267 MB** - Normalize total: 16 × 66.8 MB = **1,069 MB** - **Grand total: ~1,336 MB ≈ 1.34 GB** At 3.35 TB/s, theoretical minimum = 1.34 / 3350 = **0.40 ms** So the fused single-kernel approach should be able to reach close to 0.5ms if memory bandwidth bound. But with 3 separate kernel launches, each needs to read the full input from DRAM: 1. compute_stats: read 267 MB 2. finalize_stats: reads two tiny double arrays (negligible) 3. normalize: reads input, weight, bias = 267 + 267 + 267 = 801 MB, writes output = 267 MB Wait, the normalize kernel reads input AGAIN from DRAM. So total read = 267 + 267 + 801 = 1.335 GB, writes = 267 MB. Total memory = ~1.6 GB. Fused to single kernel: - Each block reads input once, computes stats, stores locally, then writes normalized output - Reads: input (267 MB) + weight (267 MB) + bias (267 MB) = 801 MB - Writes: output (267 MB) - Total: 1.068 GB (33% less!) This could be significantly faster! Let me write a single fused kernel. But we need to handle alignment. Let me write a cleaner version. One challenge with single-block-per-batch: 256 threads per block means 16 blocks for 16 batches. On H100 with many SMs, this might not saturate all SMs. But the memory bandwidth is the bottleneck, not compute. With 4096 threads doing coalesced reads, we should be close to memory bandwidth bound. Actually, let me think again. Each thread would use grid-stride loop. With block size 256 and one block per batch: - Thread 0 reads elements 0, 256, 512, ... across 4M elements - The pattern is strided by 256 floats = 1024 bytes = 1 KB For coalesced memory access, consecutive threads should read consecutive memory. In a grid-stride loop, within each iteration, thread i reads element i * stride + block_offset. With 256 threads, threads 0..255 in warp 0 read elements 0..255 which are perfectly coalesced. Similarly for warp 1: elements 256..511, etc. So the access pattern IS coalesced within each warp. Wait, no. Let me reconsider. In a grid-stride loop with block size 256: - Thread 0 reads index 0, then 256, then 512, etc. - Thread 1 reads index 1, then 257, then 513, etc. So within each iteration, threads 0..255 read 0..255 which is coalesced. Then threads all increment by 256. Yes! The grid-stride loop IS coalesced within each iteration. So memory access is fine. Now, the issue with single block per batch is that reduction happens within one block. With 4M elements processed by 256 threads, each thread does ~16K iterations. Wait, no: - 4M / 256 = 16,384 iterations per thread in the simple loop - With float4, that's ~4,096 iterations with 4 floats each Each iteration does 4 floats. The inner loop has overhead. This might be okay. Actually wait: 4,096,000 / (256 * 4) = 4,000 iterations of a float4 loop. That's reasonable. Hmm, but 4000 iterations per thread with ~12 arithmetic ops each = 48K ops = not compute bound. The memory reads are the bottleneck. Let me just write the fused kernel and test it. Actually, I realize there's an issue: for very large N (4M), doing a grid-stride loop with just 256 threads per block means each thread needs to iterate many times. The overhead of the loop might add up. Alternatively, I can use multiple blocks per batch: 1. First kernel: compute per-block partial sums (write to global memory) 2. Second kernel (small): reduce partial sums to get mean/rstd, then normalize Actually, this is similar to what PyTorch does internally. But writing partial sums to global memory and reading them back adds overhead. Let me think about another approach: a single kernel where each block computes stats for one batch using grid-stride, reduces within block, stores mean/rstd to shared memory, then does grid-stride normalization. But each batch is handled by MULTIPLE blocks. Wait, that won't work because only one block can write to the shared memory for mean/rstd. If multiple blocks process the same batch, they each need the same mean/rstd. Solution: 1. Pass 1 with multiple blocks per batch: compute per-warp partial sums, write to global memory 2. Pass 2 with 1 block per batch: reduce partial sums, compute mean/rstd, store to global memory 3. Pass 3 with multiple blocks per batch: normalize using mean/rstd This is essentially what PyTorch does: persistent reduction followed by element-wise transformation. But that's 3 passes again. Wait, I think a better solution is to use the warp-level approach with shared memory barrier: 1. Launch with multiple blocks per batch 2. Each block computes partial sum/sum_sq 3. Write partial results to a global scratchpad (one entry per block per batch) 4. Synchronize all blocks for a batch (using a global counter / semaphore) 5. First block to reach the barrier reduces all partial results, computes mean/rstd 6. All blocks proceed to normalize This is complex because it requires inter-block synchronization. Actually, for 16 batches and 4M elements each, I think the current 3-kernel approach should be performance-competitive. Let me optimize it instead: 1. Use more blocks per batch to increase parallelism 2. Avoid double-precision atomics by using Online algorithm (Welford) with float and careful ordering 3. But wait, float atomics might not be precise enough Actually, PyTorch uses single-precision reduction for speed and fp32 intermediate values. The small numerical error is acceptable. Let me try a different approach: Use CUDA cooperative groups for block-level reduction without shared memory roundtrips. Actually, let me try optimizing the existing kernel by: 1. Using float2 instead of float4 for more parallelism 2. Increasing blocks per batch to saturate SMs 3. Using cp.async or just ensuring good occupancy For the H100, maximum blocks per SM is 32, and we have many SMs. With batch_size=16 and blocks_per_batch, say, 64, we get 1024 blocks, which is good. But I think the simplest optimization that might help a lot is to fuse the two read passes into one. Currently the input is read twice. If I can load the input into shared memory or registers, compute stats, then reuse it for normalize... but the input is 16MB per batch which is way too big for shared memory. What if I tile the computation? Each block processes a tile of, say, 8K elements: 1. Load tile into shared memory (or registers) 2. Compute partial sum/sum_sq 3. Write partial results to global memory 4. Synchronize (across kernels) 5. Compute mean/rstd 6. Each block re-reads its tile, normalizes, writes output Still two reads of input. Can't avoid this without persistent kernel or L2 cache keeping the data (which it might!). Actually, on H100 with large L2 cache (50MB+), if the input tensors are small enough, the first read might still be in L2 when the normalize kernel reads them. For 16 batches: total input = 16 * 16.7 MB = 267 MB, which is larger than L2 cache. For a single batch: input = 16.7 MB, which might fit nicely in L2. The weight and bias are each 16.7 MB. Total working set for one batch = 50 MB, which might be close to L2 capacity. I think for this problem size, PyTorch's native implementation should already be very optimized. The key to speedup might be: 1. Better kernel fusion 2. Lower kernel launch overhead Actually wait - I should check whether PyTorch's native LayerNorm is using the full tensor or just treating the last 3 dimensions. Since normalized_shape=(64,256,256), PyTorch's LayerNorm normalizes over the last 3 dimensions, i.e., each of the 16 elements in the first dimension gets normalized independently. So it's really batch_size=16, normalized_dim=4194304. PyTorch likely uses its own highly optimized LayerNorm kernel. To beat it, I need to be clever. One thing I notice: my kernel reads weight and bias across all batches even though they're the same. Can I broadcast them more efficiently? Actually for the element-wise multiply and add, the weight and bias are read once per element. The memory bandwidth for this is fixed. The only way to reduce it is to not read them if they're trivial (all 1s and 0s), but that's not the case here since they're parameters. Let me think about another approach: use Tensor Cores! For LayerNorm, we could potentially use warp-level matrix operations. But the operation is element-wise with broadcasting, not really a matrix multiplication. OK, let me try the single fused kernel approach with grid-stride loops to see if it helps: ```cuda __global__ void fused_layer_norm_kernel( const float* __restrict__ input, const float* __restrict__ weight, const float* __restrict__ bias, float* __restrict__ output, int batch_size, int N, float eps) { // Each block handles one batch int batch_idx = blockIdx.x; if (batch_idx >= batch_size) return; const float* batch_input = input + batch_idx * N; float* batch_output = output + batch_idx * N; float sum = 0.0f; float sum_sq = 0.0f; // Grid-stride loop for stats for (int i = threadIdx.x; i < N; i += blockDim.x) { float val = batch_input[i]; sum += val; sum_sq += val * val; } // Block reduce __shared__ float s_sum[256]; __shared__ float s_sum_sq[256]; __shared__ float s_mean; __shared__ float s_rstd; s_sum[threadIdx.x] = sum; s_sum_sq[threadIdx.x] = sum_sq; __syncthreads(); // Reduce for (int stride = blockDim.x / 2; stride > 0; stride /= 2) { if (threadIdx.x < stride) { s_sum[threadIdx.x] += s_sum[threadIdx.x + stride]; s_sum_sq[threadIdx.x] += s_sum_sq[threadIdx.x + stride]; } __syncthreads(); } if (threadIdx.x == 0) { float mean = s_sum[0] / N; float var = s_sum_sq[0] / N - mean * mean; s_mean = mean; s_rstd = rsqrtf(var + eps); } __syncthreads(); float m = s_mean; float rs = s_rstd; // Grid-stride loop for normalize for (int i = threadIdx.x; i < N; i += blockDim.x) { float val = batch_input[i]; batch_output[i] = (val - m) * rs * weight[i] + bias[i]; } } ``` Wait, there's a bug: `weight[i]` should be `weight[i]` where weight has shape [N], but it's actually [features, dim1, dim2]. Since N = features * dim1 * dim2, and weight is contiguous, `weight[i]` is correct. But there's another issue: with 256 threads per block and N=4M, the grid-stride loop does 16K iterations per thread! That will take forever due to loop overhead. Let me use float4 to reduce iterations: ```cuda __global__ void fused_layer_norm_kernel(...) { // ... int num_float4 = N / 4; // Stats pass float4 sum4 = make_float4(0,0,0,0); float4 sum_sq4 = make_float4(0,0,0,0); for (int i = threadIdx.x; i < num_float4; i += blockDim.x) { float4 val = reinterpret_cast<const float4*>(batch_input)[i]; sum4.x += val.x; sum_sq4.x += val.x * val.x; // ... etc } float sum = sum4.x + sum4.y + sum4.z + sum4.w; float sum_sq = sum_sq4.x + sum_sq4.y + sum_sq4.z + sum_sq4.w; // remainder for (int i = num_float4 * 4 + threadIdx.x; i < N; i += blockDim.x) { float val = batch_input[i]; sum += val; sum_sq += val * val; } // block reduce... } ``` This reduces iterations from 16K to 4K per thread. Still a lot but manageable. Actually, to really saturate bandwidth, I should use MORE threads per batch. Let me use 512 or 1024 threads per block with multiple blocks per batch. Hmm, let me think about this differently. For LayerNorm, the key optimization opportunities are: 1. Fusing the reduction with the normalize (to reuse input from cache/L2) 2. Vectorized loads (float4) 3. Reducing block-level synchronization overhead For a sophisticated implementation, I should look at what PyTorch does. PyTorch's LayerNorm likely uses: - cuDNN (if available) which is highly optimized - Or a custom Triton kernel - Or the PyTorch ATen implementation Since cuDNN might not be available or might not handle this specific shape optimally, PyTorch might fall back to its own kernel. Let me try a different strategy: use persistent warps that tile the problem at the warp level, and use warp-shuffle for reduction. Actually, I just realized: the simplest optimization I haven't tried is to make the number of blocks per batch much larger. Currently I use `blocks_per_batch = min(1024, ...)`. Let me make it even larger, like 256 or 512 blocks per batch (4096 or 8192 total blocks), which should saturate the GPU better. Wait, but this requires more global memory for storing per-block partial results. Let me rewrite to use a two-stage reduction: 1. First kernel: many blocks per batch compute partial sums, store per-block results 2. Second kernel: small kernel reduces per-block partials, computes mean/rstd 3. Third kernel: many blocks normalize This is cleaner and avoids atomics. Let me write this version. Actually, thinking more carefully: - With blocks_per_batch = 256, we have 4096 blocks total for 16 batches - Each block has 256 threads, so ~1M threads total - But the memory bandwidth limits us, not thread count For compute_stats: we only need ~267MB of reads. At 3.35 TB/s, this is ~0.08ms. But kernel launch and reduction overhead add more. For normalize: we read input (267MB) + weight (267MB) + bias (267MB) = 801MB and write output (267MB). At 3.35 TB/s, ~0.32ms. So total theoretical minimum is ~0.4ms + overheads. My current kernel likely does about 0.6-1ms. PyTorch's LayerNorm probably also does around 0.5-0.8ms. To beat it, I need to reduce overhead and maybe fuse kernels. Let me go with the single fused kernel approach but use MORE blocks per batch. For the stats phase, each block computes partial sums. Then I need a way to reduce these partial sums and broadcast mean/rstd to all blocks. One way: use global memory atomic counter for synchronization between blocks of the same batch. This is tricky but can work. Simpler: do two kernel launches. 1. First: compute stats (mean, rstd) - one kernel 2. Second: normalize using computed stats - one kernel This is what I already have, minus the separate finalize kernel. Wait, I have THREE kernels: compute_stats, finalize_stats, normalize. I can merge compute_stats and finalize_stats. Let me try merging compute_stats and finalize_stats into a single kernel that does full reduction: - Use many blocks per batch for the first pass, each computing partial sums - Store partial sums in a global scratchpad - Last block for each batch reduces all partial sums and writes mean/rstd For this I need to know how many blocks per batch. Let me set it to, say, 64 blocks per batch. ```cuda __global__ void compute_stats_fused_kernel(...) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; int tid = threadIdx.x; // Compute partial sum and sum_sq // ... // Store to global scratchpad int scratch_offset = batch_idx * blocks_per_batch + block_in_batch; partial_sum[scratch_offset] = block_sum; partial_sum_sq[scratch_offset] = block_sum_sq; // Synchronize within batch using atomic counter __threadfence(); int old = atomicAdd(&counter[batch_idx], 1); if (old == blocks_per_batch - 1) { // Last block: reduce all partial sums double total_sum = 0, total_sum_sq = 0; for (int i = 0; i < blocks_per_batch; i++) { total_sum += partial_sum[batch_idx * blocks_per_batch + i]; total_sum_sq += partial_sum_sq[batch_idx * blocks_per_batch + i]; } mean[batch_idx] = total_sum / N; rstd[batch_idx] = rsqrt(total_sum_sq / N - mean[batch_idx] * mean[batch_idx] + eps); } } ``` The inter-block sync via atomic is tricky because we need to ensure all writes to partial_sum are visible before reading. `__threadfence()` helps with this but there's still a race condition window. For correctness in production, we'd want a more robust approach. Actually, for 16 batches and a small number of blocks per batch, I can do the reduction on the CPU or with a very small second kernel. Two kernel launches is fine. Let me go with a simpler two-kernel approach: 1. Kernel 1: Many blocks per batch compute partial sums, write to global memory 2. Kernel 2: One block per batch reduces partial sums to mean/rstd, then all blocks normalize Wait, this still requires a third kernel for normalize if we want many blocks. OK, best practical approach: 1. Kernel 1: Compute partial sums with many blocks, write to global memory 2. Kernel 2: Small kernel reduces partial sums per batch to mean/rstd 3. Kernel 3: Normalize with many blocks This is 3 kernels. Same as my current approach. But I can avoid atomics in kernel 1 by using global memory writes instead. The key difference: my current approach uses atomicAdd in the reduction phase. For 64-256 blocks contending on 16 double-precision atomic locations, this might be slow. Let me switch to the global memory approach for the reduction. Actually, I just realized: my current kernel has only 2 atomic operations per block (one for sum, one for sum_sq). With, say, 64 blocks per batch × 16 batches = 1024 blocks, and each block writing 2 atomics, that's 2048 atomic operations total. These are spread across 16 memory locations (one per batch), so each location gets 128 atomic adds. Double-precision atomics are slower than single-precision. A better approach: have each block write its partial sum to a unique global memory location, then reduce those in a separate kernel. Let me implement this two-stage reduction. Actually, for simplicity and to test if it matters, let me just try increasing the block size and see if performance improves. Wait, I should also try a completely different approach: using block-level persistence where each block stays resident and processes elements while keeping them in registers/shared memory between stats and normalize. But this requires single-block-per-batch which limits parallelism. For 16 batches with a block size of 512 or 1024, we'd have 16 blocks of 1024 threads = 16,384 threads total. H100 has ~132 SMs, so we're using a small fraction. But each block needs to read 4M elements × 4 bytes = 16MB. At 3.35 TB/s with 16 parallel reads... the bandwidth is shared. With 16 blocks each reading 16MB sequentially, the total bandwidth demand is 16 × 16MB / (time per block). If each block takes T seconds, total reads = 256MB / T. To use full bandwidth: 256MB / T = 3350 GB/s → T = 0.076ms. During this time, each block processes 16MB in 0.076ms, which is 210 GB/s per block. The H100 can handle much more per SM. So we'd be under-utilizing memory bandwidth. With more blocks per batch, say 64 blocks per batch × 16 batches = 1024 blocks, each reading 4M/64 = 65K elements = 256KB. Total parallelism is much higher. OK, let me take the current 3-kernel design but optimize the reduction. Actually, let me stick with the current design and just try to optimize a few things: 1. Use float instead of double for stats to match PyTorch (which uses float in its fast path) 2. Use warp-shuffle reduction more efficiently 3. Use more blocks per batch Actually, I worry that using float instead of double might cause numerical issues. Let me check what PyTorch does for fp32 layer norm. PyTorch's native LayerNorm uses float32 accumulation for the reduction in its fast kernels (the Triton-based ones in PyTorch 2.0+). Let me try with single-precision atomics to a float accumulator. But wait, atomicAdd for float32 exists and is faster. Hmm, but the reference uses nn.LayerNorm which might use fp32 internally even for accumulation. Let me change to float accumulation and see if correctness still passes. The main numerical concern is the E[x^2] - E[x]^2 formula for variance. For large N (4M), the sums can be large and subtracting two large numbers loses precision. Welford's online algorithm is better for numerical stability. Actually, PyTorch's LayerNorm definitely uses Welford's algorithm or a numerically stable two-pass approach with compensated summation. Let me implement Welford's algorithm per thread, then combine across threads. But for matching the reference, as long as we're within fp32 tolerance (1e-4), either approach should work. OK, let me try a new design: 1. Keep the 3-kernel structure 2. Use single-precision float for accumulation 3. Use Welford's online algorithm per-thread 4. Warp-shuffle reduction 5. Block-level reduction with shared memory 6. Single atomic per block that writes a struct of (mean, m2, count) or just use global scratchpad Wait, I think the simplest improvement is to just avoid the separate finalize kernel by handling the final reduction differently, and to use Welford's algorithm for numerical stability. Actually, I notice my current design works and passes correctness. Let me try to optimize it by: 1. Removing the finalize_stats kernel and doing the final computation in the first kernel using a better reduction strategy 2. Using more blocks per batch for better parallelism Let me write a version where kernel 1 computes partial sums with many blocks per batch, kernel 2 reduces the partial sums and normalizes. Actually wait, the simplest change might be: just increase `elements_per_thread` in the first kernel so that `blocks_per_batch` is smaller, reducing atomic contention. Currently with `elements_per_thread=64` and `N=4M`, `blocks_per_batch ~ 256`. With `elements_per_thread=256`, `blocks_per_batch ~ 64`. Fewer blocks = less atomic contention. Or conversely, if I want more parallelism, I can reduce `elements_per_thread` but use the scratchpad approach instead of atomics. Let me just try the scratchpad approach with many blocks. First kernel writes partial results, second kernel reduces and normalizes. Actually, here's the most practical improvement: I'll write a two-kernel solution where the first kernel computes per-block partial (sum, sum_sq), and the second kernel reduces those partials and normalizes. Wait, but the second kernel needs to handle both reduction (which only needs 1 block per batch) and normalization (which needs many blocks). Unless I launch two separate kernels. Actually, I can do it with 2 kernels: 1. compute_stats: many blocks per batch, write partial sums to scratchpad 2. normalize: many blocks per batch, first warp of first block reduces partials, all blocks normalize But doing reduction inside normalize requires that the first block finishes reduction before others start reading mean/rstd. This requires synchronization. Actually, let me just try to optimize the current 3-kernel design. The key is: 1. Make sure the number of blocks is optimal 2. Use vectorized memory access 3. Minimize register usage for better occupancy For compute_stats, with `blocks_per_batch = 64`, we have 1024 blocks total. Each block reads 4M/64 = 65,536 elements. At 256 threads per block, each thread reads 256 elements, which is 64 float4 loads. This is reasonable. Let me update my kernel with these optimizations: - Reduce elements_per_thread to get blocks_per_batch = 64 or 128 - Make sure float4 alignment is properly handled - Remove the separate finalize kernel by using a different approach Actually, let me look at what happens if I just fuse the finalize_stats into compute_stats. Can I have the first block of each batch do the final reduction after all blocks have written their partial results? With CUDA, inter-block synchronization within a single kernel launch is tricky. One approach: - Use an atomic counter per batch - Last block to arrive does the reduction and writes mean/rstd - All blocks then spin-wait until mean/rstd is ready - Then all blocks normalize This is the persistent/kernel-fusion approach. Let me write it. ```cuda __global__ void fused_stats_normalize_kernel(...) { int batch_idx = blockIdx.x / blocks_per_batch; int block_in_batch = blockIdx.x % blocks_per_batch; // ... // Compute partial sum and sum_sq // Store to scratchpad[block_in_batch] __threadfence(); // Atomic counter __shared__ bool is_last_block; if (threadIdx.x == 0) { int old = atomicAdd(&counter[batch_idx], 1); is_last_block = (old == blocks_per_batch - 1); } __syncthreads(); if (is_last_block) { // Reduce all partials // Write mean[batch_idx], rstd[batch_idx] // Reset counter for next launch } // Barrier: all blocks wait until mean/rstd is ready // This requires spinning on a variable, which is bad for GPU scheduling } ``` This has performance issues because spinning blocks waste SMs. A better approach: use two separate kernel launches. First kernel computes stats. Second kernel does normalize. This is clean and can be fast if L2 cache keeps the input between kernels. For this problem size (16 batches × 4M elements = 256MB total input), the L2 cache on H100 is 50MB, so only a fraction stays cached. But the weight and bias (16.7MB each) fit in L2, and each batch's input (16.7MB) might stay if the kernels are launched back-to-back. OK, let me just write a clean 2-kernel fused version: 1. First kernel: many blocks per batch compute partial sums using Welford, write to scratchpad 2. Second kernel: small grid (1 block per batch) reduces scratchpad, then large grid normalizes Wait, that requires global barriers between reduction and normalize, which means 2 kernel launches anyway. So: 2 kernels total. 1. compute_stats: many blocks, write partial results 2. reduce_and_normalize: first do block-level reduction of partials, then normalize For kernel 2, I can have it first reduce partials and then normalize in the same kernel. But normalization needs many blocks for parallelism. Unless I oversubscribe: launch enough blocks to saturate SMs, and have each block first do a small piece of reduction (if it's one of the first blocks_per_batch blocks), then all blocks do normalization. This is still messy. Let me just go with the 3-kernel approach but optimize it. Actually, I realize the cleanest solution is: **Kernel 1 (compute_stats)**: Many blocks per batch compute partial sum and sum_sq. Write to global scratchpad (2 × blocks_per_batch × batch_size floats). No atomics needed. **Kernel 2 (finalize)**: One block per batch reads the partial sums, reduces them, writes mean and rstd. **Kernel 3 (normalize)**: Many blocks per batch normalize using precomputed mean and rstd. The key benefit over my current approach: kernel 1 uses NO ATOMICS. It just writes to unique global memory locations. This is much faster than atomicAdd, especially for double-precision. Let me implement this. Each block writes to `scratch[batch_idx * blocks_per_batch + block_in_batch]` for sum and similarly for sum_sq. For kernel 2, with blocks_per_batch = 64, each block reads 64 floats and 64 floats, reduces them. This is trivial. Kernel 3 is the same as before. Let me write this version. For scratchpad memory management, I need to allocate it dynamically or use a fixed-size array. I'll allocate it dynamically in the C++ function. The total scratchpad size = 2 × batch_size × blocks_per_batch × sizeof(double) = 2 × 16 × 64 × 8 = 16KB. Very small. Actually, I can use float instead of double for the scratchpad to save memory and bandwidth. The partial sums from each block won't overflow float32 for 65K elements with values in [0,1]. Max partial sum: 65,536 × 1.0 = 65,536, which fits in float32 with plenty of room. Max partial sum_sq: 65,536 × 1.0 = 65,536, same. And the final reduction sums 64 partials: max total = 64 × 65,536 = 4,194,304 = N. Sum_sq max = same. Both fit in float32 precisely (4M is well within 24-bit mantissa precision). So float32 scratchpad is fine and much faster. Let me write this optimized version. Wait, but PyTorch's LayerNorm might use Welford's algorithm, which computes running mean and M2. Two-pass E[x^2] - E[x]^2 can have precision issues for large variance or when mean is large. But for input ~[0,1], the mean is around 0.5 and variance is around 0.08. The difference E[x^2] - E[x]^2 = 0.33 - 0.25 = 0.08. This is well-conditioned for float32 with N=4M. Actually there's a subtle numerical issue: computing variance as E[x^2] - E[x]^2 can give negative results due to roundoff, which is why we typically clamp to 0 or use a small eps. With float accumulation and N=4M, roundoff in the sum could be on the order of N * eps_machine * typical_value ≈ 4M * 1e-7 * 1 = 0.4. So sum might have error ~0.4. Then E[x] = sum/N has error ~0.4/4M = 1e-7. E[x^2] similarly. The variance = E[x^2] - E[x]^2 could have error up to ~2e-7, which is tiny compared to typical variance of 0.08. With eps=1e-5, we're fine. But wait, the atomic add to double accumulates multiple partial sums. The final double-accurate sum divided by N gives a very accurate mean. With float scratchpad, the partial sums are accurate to ~1e-7 relative (since each partial sum is 65K values, the roundoff is ~65K * 1e-7 * 1 = 0.0065). Summing 64 partial sums: error could be ~64 * 0.0065 = 0.42 in the worst case (much worse than the double atomic approach which had ~0.4 total error). So actually, the double atomic approach might be MORE accurate for the mean than float scratchpad! But the mean with float scratchpad is still accurate enough for our tolerance (atol=1e-4). OK, I'll use float scratchpad for speed and see if correctness still passes. Actually, let me think more carefully. My current double-atomic approach: - Each block computes precise double sum of its elements - AtomicAdd adds these doubles to global double sum - Final double sum is accurate to ~machine epsilon of double (~1e-16) - Mean = sum/N is extremely accurate Float scratchpad approach: - Each block computes float sum of its elements (error ~0.0065 per block for 65K elements) - We sum 64 floats (error could accumulate to ~0.4) - Mean = sum/N has error ~0.4/4M = 1e-7 So mean error is ~1e-7 in both cases. The double approach gives ~1e-13 mean error. Both are well within fp32 tolerance of 1e-4. For variance: Double approach: E[x^2] computed with double precision, extremely accurate. Float scratchpad: same ~1e-7 relative error. The variance is then computed as E[x^2] - E[x]^2. With double, this difference is computed in double precision then cast to float. With float, the difference might lose precision if the values are close. For our data (uniform [0,1]), E[x^2]≈0.33, E[x]^2≈0.25, difference≈0.08. The relative error in the difference depends on the absolute errors. With mean error ~1e-7 and E[x^2] error ~1e-7, the variance error is ~(1e-7 + 1e-7) = 2e-7. Relative to 0.08, that's ~2.5e-6. Well within tolerance. So float scratchpad should be fine numerically. Let me write the improved version now. Actually, I realize I can simplify further: if I use warp-level reduction with warp-shuffle and store per-warp results, I can reduce the scratchpad size. But that adds complexity. Let me just use per-block partial sums. Here's my plan: 1. Kernel 1: `compute_partial_stats` - Each block computes sum and sum_sq for its chunk of elements - Writes to scratchpad[batch_idx * blocks_per_batch + block_in_batch] for sum and sum_sq 2. Kernel 2: `reduce_and_normalize` - Wait, I need mean and rstd for normalization. Let me do: Actually, let me do 2 kernels: 1. `compute_stats`: compute partial sums, write to scratchpad. Use atomic counter to track completion. 2. `normalize`: use precomputed mean and rstd. Hmm, I still need a separate reduction step. Let me do 3 kernels but with the first two merged if possible. Wait, I just realized I can merge the reduction and normalization in a clever way: Kernel 1: compute partial stats → scratchpad Kernel 2: - Phase A: Each block's warp 0 reduces a portion of the scratchpad - Phase B: Block 0 of each batch finishes the reduction, computes mean/rstd, stores to output array - Phase C: All blocks normalize But Phase C depends on Phase B completing, which requires synchronization across blocks of the same batch. This is the inter-block sync problem again. OK, simplest clean approach: 3 kernels, no atomics in kernel 1. Actually you know what, let me just make the first kernel use float scratchpad and see if it's faster. I'll write it cleanly. Wait, I just had a better idea. The final kernel (`normalize`) currently reads the entire input from DRAM again. But if I launch the 3 kernels back-to-back, the L2 cache might retain some of the input from the first kernel's read. For each batch (16MB of input), the L2 cache on H100 is 50MB, so the entire batch's input PLUS weight (16MB) could stay in L2 if we process one batch at a time. But we process all 16 batches in parallel with many blocks. Actually, there's no guarantee of L2 caching across kernels, so I should assume DRAM reads for each kernel. Let me calculate: for kernel 1, we read 267MB. For kernel 3, we read input (267MB) + weight (267MB) + bias (267MB) = 801MB and write output (267MB). Total memory traffic = 267 + 801 + 267 = 1,335MB. If I fuse kernels 2 and 3 (reduce and normalize in one kernel), the traffic is the same, but I save one kernel launch. If I somehow fuse ALL 3 into one kernel with inter-block sync, I'd save two kernel launches. But kernel launches are only ~5-10 microseconds, negligible compared to 0.5ms compute time. Let me just write the clean float-scratchpad version and see. Actually, I realize I should profile before optimizing further. But since I can't profile directly, let me run correctness again with a slightly optimized version and then submit. Hmm wait, I can run correctness to verify, but I don't get timing info. Let me just write what I think is the best version and submit it. Let me write a cleaner version with: 1. No atomics in stats computation - use scratchpad 2. Two-stage reduction using warp shuffles 3. Single-precision throughout for speed (matching PyTorch's fast path) Actually, for numerical correctness, let me keep double for the stats accumulation to stay close to the reference's precision. Here's my refined approach: 1. Kernel 1: Many blocks per batch. Each block computes sum and sum_sq in double precision using warp shuffles. Writes per-block results to double scratchpad. 2. Kernel 2: One block per batch. Reads scratchpad, reduces with warp shuffles, computes mean and rstd in double precision, writes to float output arrays. 3. Kernel 3: Many blocks per batch. Normalizes input using mean/rstd, weight, bias. The key improvement: kernel 1 has NO ATOMICS. It just writes to unique global memory locations. This is much faster, especially for double-precision writes. Let me implement this.