reasoning
Let me analyze this problem:
1. We need to implement a fused kernel for: GEMM (Linear) -> GroupNorm -> Swish -> Multiply -> Swish
2. The GPU is H100 with 3.35 TB/s memory bandwidth and very high compute (67 TFLOPS FP32, 1979 TFLOPS FP16/BF16 Tensor Core)
3. Input: batch_size=1024, in_features=8192, out_features=8192, num_groups=256, multiply_weight_shape=(8192,)
Let me think about the operations:
- Linear: (1024, 8192) x (8192, 8192) = very large GEMM, ~134M MACs
- GroupNorm: 256 groups on 8192 features, so 32 features per group. For batch_size=1024.
- Swish: x * sigmoid(x)
- Multiply: x * weight
- Swish: x * sigmoid(x)
The key insight is that we can fuse everything after the GEMM into a single kernel. But actually, we can't easily fuse the GEMM with GroupNorm because we need the full GEMM output to compute GroupNorm statistics.
Wait, we could use cuBLAS for the GEMM and then fuse the rest. But given the H100 has amazing Tensor Cores, we might want to use FP16/BF16 Tensor Cores if possible... but the requirement says "fp32 precision" for correctness, with atol=rtol=1e-4.
Actually, let me re-read: "Produce numerically equivalent outputs at fp32 precision." This means the output should match the reference model which runs in fp32. But we can still compute internally using Tensor Cores and convert back, as long as the numerical difference is within 1e-4.
However, for simplicity and guaranteed correctness, let me first try keeping GEMM in FP32 using cuBLAS, and fusing the rest.
Actually, a key optimization: Linear in PyTorch already uses highly optimized cuBLAS. The GroupNorm, then two Swish and elementwise multiply are memory-bound. Fusing them into one kernel would save multiple passes over the data.
For 1024 x 8192 = ~8.4M elements:
- GroupNorm reads once, writes once
- Swish reads once, writes once
- Multiply reads once, writes once
- Swish reads once, writes once
Total: 5 reads + 4 writes for the non-GEMM part if unfused.
Fused: 1 read + 1 write.
But actually, can we do better? On H100, maybe we can fuse the whole thing including GEMM using warp-specialization or similar?
Let me think about this differently. The linear layer produces output. For GroupNorm, we need mean and variance per group. num_groups=256, out_features=8192, so 32 features per group. batch_size=1024.
For each group, we process batch_size * features_per_group = 1024 * 32 = 32768 elements.
We can compute mean and variance in shared memory. Since there are 256 groups, and for each sample in the batch we have one value per feature...
Actually, GroupNorm computes statistics across both batch and features within each group. Wait no - let me re-check:
PyTorch GroupNorm:
- Input: (N, C) or (N, C, *)
- Divides C into num_groups groups
- Computes mean and std over (N, *) for each group
- Normalizes
So for input (1024, 8192) with num_groups=256:
- Each group has 8192/256 = 32 channels
- Mean and variance are computed over (batch_size, 32) = (1024, 32) = 32768 elements per group
- There are 256 groups
So we can have each thread block handle one or more groups. With 32768 elements per group, that's quite a lot.
For the GEMM, since it's (1024, 8192) @ (8192, 8192), this is a large GEMM that will be compute-bound. We should definitely use cuBLAS for this.
Strategy:
1. Use `torch.matmul` or the linear layer weight for GEMM
2. Write a fused kernel for GroupNorm + Swish + Mul + Swish
Can we avoid the extra memory write/read between GEMM and the fused kernel? Not easily unless we write our own GEMM, which is probably not worth it given cuBLAS is very optimized.
Wait, but maybe we can use CUTLASS or a custom GEMM kernel that directly writes to the computation... but writing a competitive GEMM kernel for (1024, 8192) @ (8192, 8192) on H100 is extremely hard.
Actually, let me check: could we use the highly optimized Linear and then do a custom fused kernel for the rest?
Let me calculate:
- GEMM: 1024 * 8192 * 8192 * 2 = ~137 GFLOPs. At 67 TFLOPS, this takes ~2ms in theory.
- Memory bandwidth for the rest: each element is 4 bytes in fp32.
- Fused kernel: read 8.4M * 4 = 33.6 MB, write 33.6 MB = 67.2 MB total
- At 3.35 TB/s: ~20 microseconds
- Unfused (separate ops): much more memory traffic
Actually, the PyTorch reference would have:
1. Linear: compute-bound, fast cuBLAS
2. GroupNorm: reads 33.6MB, writes 33.6MB, plus reduction operations
3. Swish: reads 33.6MB, writes 33.6MB
4. Mul: reads 33.6MB (data) + 33.6KB (weight), writes 33.6MB
5. Swish: reads 33.6MB, writes 33.6MB
Total memory for post-GEMM: ~335MB + reductions.
If we fuse them all: ~67MB + reductions.
So fused kernel could save significant time.
But wait - can we do even more? What about the GEMM output being consumed immediately? The issue is that GroupNorm needs reductions across batch elements, so we need the full GEMM output first. However, since the batch_size is only 1024 and out_features=8192, the output is just 32MB. We could potentially tile the GEMM and process GroupNorm on tiles... but that's complex.
Let me think about a custom kernel design for the post-GEMM operations.
Kernel design:
- Each thread block handles one group, processing all batch elements
- Actually, better: each thread block handles multiple channel elements for all batch items
- Since we need per-group mean/variance, we need a reduction within each group first
For GroupNorm with 256 groups and 1024 batch size:
- We can launch 256 thread blocks, one per group
- Each thread block processes 1024 * 32 = 32768 elements
- Use shared memory for reduction
But wait, after GroupNorm we need:
`x = x * sigmoid(x) * weight * sigmoid(x * sigmoid(x) * weight)`
Actually let me re-read the ops:
```python
x = x * torch.sigmoid(x) # Swish
x = x * self.multiply_weight # Multiply
x = x * torch.sigmoid(x) # Swish
```
So:
`output = (x * sigmoid(x) * weight) * sigmoid(x * sigmoid(x) * weight)`
This is what we need to compute element-wise after GroupNorm.
Kernel strategy:
1. Compute GEMM using standard Linear
2. Fused kernel:
- For each group, compute mean and variance across (batch_size, channels_per_group)
- Then for each element, apply normalization, then the two Swish and multiply
Let me think about the kernel more carefully.
For step 1 (mean/variance):
- Grid: 256 blocks (one per group)
- Each block: processes all 1024 * 32 elements of its group
- Use shared memory reduction to compute mean and inv_std
Actually, we need to be careful about numerical precision. We can use Welford's online algorithm for better numerical stability.
Channels per group = 8192 / 256 = 32.
For each block (group g):
- Elements are at positions [0..1023][g*32..(g+1)*32-1] in row-major (batch, out_features)
- But in PyTorch, Linear output is (batch_size, out_features) which is row-major.
So for group g, elements are:
- batch i, channel c where g*32 <= c < (g+1)*32
In memory: index = i * 8192 + c
For reduction, we need to compute mean and variance of all elements where c is in the group's range.
Launch config:
- 256 blocks (for 256 groups)
- Each block: maybe 256 or 512 threads
- Each thread processes 32768 / num_threads elements
After computing mean and inv_std in shared memory, we need to write the normalized values and apply the rest of the ops. But wait - if we have each block do the reduction first, then write, we'd need a second pass or we could interleave.
Actually, since we need mean and inv_std for ALL elements in the group before normalizing, we have two natural approaches:
1. Two-pass: first compute stats, store to shared memory, __syncthreads(), then each thread reads its elements, normalizes, and applies the remaining ops
2. Stream data through registers if the block can hold all data (32K elements * 4 bytes = 128KB, too big for smem)
So two-pass it is, but with coalesced memory access.
Better approach: Let me think about tiling differently. What if we process the linear + groupnorm fusion with a different tiling where we don't need to read all elements twice from global memory?
Actually, since we're reading from global memory anyway, and H100 has 3.35 TB/s, let's compute:
- Read 32MB for GroupNorm reduction: ~10 microseconds
- Write 32MB after fused ops: ~10 microseconds
The reduction overhead is additional compute but should be small.
Wait, I have an idea. What if we compute the linear layer using a custom kernel too? For a (1024, 8192) @ (8192, 8192) GEMM...
Actually, let's check what PyTorch's cuBLAS does for this. On H100, a 1024x8192x8192 GEMM should be very fast, probably under 1ms if using Tensor Cores with FP16/BF16, or maybe 1-2ms with FP32.
But the problem is that after the GEMM, PyTorch would write to global memory, and then our kernel would read it. That's one extra pass over memory (33.6MB read + write).
Could we win by doing the GEMM in a custom kernel and immediately starting GroupNorm? We'd need to tile the GEMM so that output tiles are kept in shared memory while doing the GroupNorm. But GroupNorm requires all batch elements for each channel group... which means we need the entire output before GroupNorm. Unless we do multiple tiles and accumulate mean/variance incrementally.
Yes! We can compute mean and variance incrementally using Welford's algorithm. If we tile the GEMM by output channel groups, we can:
1. Compute a tile of the GEMM output
2. Update running mean/variance for those channel groups
3. After all tiles are done, finalize mean/variance
4. Then do another pass over the GEMM output (or compute it again) to apply normalization and elementwise ops
But computing the GEMM twice is wasteful. So we'd need to store the GEMM output somewhere. The output is only 32MB, which is small. The savings from avoiding one 32MB memory pass might not be worth the complexity of writing a custom GEMM kernel.
Hmm, actually wait. If we just use PyTorch matmul and then a fused kernel, the total memory for that 32MB read is ~10 microseconds. A cuBLAS GEMM of this size on H100 is probably around 0.5-1ms. So we're talking about a 1% optimization at best by trying to fuse GEMM output into smem. Not worth it.
Let me focus on:
1. Use `torch.nn.functional.linear` or `torch.matmul` for the GEMM
2. Fused kernel for GroupNorm + 2xSwish + Multiply
Now, how to write the fused kernel efficiently?
For GroupNorm on (1024, 8192) with 256 groups:
- Each group has 32 channels
- Statistics are computed over 1024 * 32 = 32768 elements
Thread block design:
Option A: 256 blocks (one per group), each with 256 threads
- Each thread handles 32768/256 = 128 elements
- For computing stats: each thread reads 128 elements strided
- For writing back: each thread writes 128 elements
Wait, but elements in a group are not contiguous in memory! In row-major layout:
- batch 0: channels [0..31] are group 0, [32..63] are group 1, etc.
- batch 1: same pattern
So for group g, elements are at: i * 8192 + g*32 + c where 0 <= i < 1024, 0 <= c < 32.
If we use 256 threads per block, each with tid, we can have:
- Thread tid processes elements for specific channels and batches
- For coalesced access, threads should access contiguous memory
Actually, if 32 threads access 32 contiguous channels for the same batch item, that's coalesced (32 * 4 = 128 bytes, perfect for H100).
Let me redesign:
- Block = group g
- Warp = process batch items in groups
- Actually: We have 1024 batch items, each with 32 channels for our group.
- Total of 32768 elements.
For coalesced reading:
- Thread i in warp can read channel c = i % 32 for batch b = something
- Better: 32 threads read 32 channels of one batch item → coalesced
- Then move to next batch item
So if we have 256 threads = 8 warps:
- Each warp processes 32 channels of different batch items, or we interleave
- Actually, 256 threads, if we have them read 32 channels at a time: each "iteration" handles 8 batch items (since 256/32 = 8 warps, but each warp has 32 threads)
- Wait, 1 warp = 32 threads. If all 32 threads read channels 0..31 of the SAME batch item, that's coalesced. Then they move to the next batch item. So 1 warp processes 32 batch items by having each thread handle one batch item but read all 32 channels? No, that would be strided and uncoalesced.
Let me reconsider. For coalesced access on (batch, channel):
Memory layout: [batch0_ch0, batch0_ch1, ..., batch0_ch8191, batch1_ch0, ...]
For group g (channels g*32 to g*32+31):
Elements: batch0_ch(g*32) to batch0_ch(g*32+31) are CONTIGUOUS (32 floats = 128 bytes)
batch1_ch(g*32) to batch1_ch(g*32+31) are also contiguous, but separated by 8192 floats.
So: For each batch item, the 32 channels of a group are contiguous in memory.
Strategy for coalesced access:
- Each thread reads one channel from one batch item, or
- Warp reads all 32 channels of one batch item (coalesced), then moves to next batch item
With 256 threads per block:
- We could have 8 warps
- Each warp handles 1024/8 = 128 batch items
- But for perfect coalescing, each warp reads 32 contiguous floats per batch item
- Actually, we can have each thread in a warp handle one of the 32 channels, and all threads read the same batch item simultaneously → coalesced 128-byte read
Then each thread would process 1024 batch items, but only for 1 channel.
- 256 threads = 8 warps * 32 threads
- If each thread handles 1 channel, we can have 8 groups of 32 threads, each group handling 1 channel across all batch items
- But that's 256 channels, and we only have 32 channels per group...
Better:
- Use 32 threads per block (1 warp), each handling one channel
- Each thread reads all 1024 batch items for its channel, strided by 8192
- This is NOT coalesced within a warp! Thread 0 reads [0, 8192, 16384, ...], thread 1 reads [1, 8193, 16385, ...]. These are 4 bytes apart initially but then 32KB apart. This is actually NOT coalesced for the first read.
Actually, let me reconsider the memory access pattern.
If thread i reads address base + i * stride, for the first read (batch 0), threads 0..31 read addresses offset by 0, 4, 8, ... 124 bytes. This IS a coalesced 128-byte transaction! But then thread i moves to base + 8192*4 + i*4, which is another coalesced 128-byte transaction (for batch 1).
So yes, having 32 threads where thread i reads channel g*32+i for all batch items IS coalesced!
With 32 threads per block (1 warp):
- Each thread processes 1024 elements (one channel across all batches)
- Total elements processed: 32 * 1024 = 32768 ✓
But 32 threads is too few to fill the GPU. We need more blocks and more warps.
Alternative: 256 threads = 8 warps.
- But there are only 32 channels. So we can have:
- 8 warps, each warp handles 4 channels? (32 channels / 8 warps = 4 channels per warp, 4 threads per channel)
- Or: each warp handles 32 channels but only 1024/8 = 128 batch items per warp
Option: 8 warps, each warp processes 128 batch items for all 32 channels.
- In warp w, thread t (0..31):
- Handles channel c = t % 32 for batch b = w * 128 + t / 32? No, 32 threads, 32 channels.
- Actually, each thread in warp reads all 32 channels for its assigned batch items, but that would be 32 reads per batch item per thread, not coalesced if by channel.
Actually, best approach: Use 32 threads (1 warp) per block, but launch 256 blocks. But then we have only 1 warp per SM, bad occupancy. We can also have multiple blocks per SM.
Wait, 256 groups, 1024 batch items. We can have:
- 256 blocks (one per group), each with multiple warps
- Each block has, say, 8 warps = 256 threads
- The 8 warps work on different subsets of batch items: warp 0 batches 0..127, warp 1 batches 128..255, etc.
- Within each warp, 32 threads handle 32 channels, one channel per thread
- Thread i in warp w: processes channel (group*32 + i) for batch items [w*128, w*128+127]
This gives us:
- Coalesced reads: Each warp reads 128 batch items. For batch b, thread i reads address b*8192 + group*32 + i. This is coalesced for the 32 threads reading batch b simultaneously.
- But different warps read different batch items, so no conflicts.
Wait, for coalescing in a warp, all 32 threads should access consecutive memory addresses. If thread i reads address base + i for the same batch item, and we iterate over batch items, then:
- Iteration b: thread i reads b * 8192 + group * 32 + i. These addresses differ by 4 bytes, so they span 128 bytes. Perfect coalescing!
Yes! This works. With 8 warps and 128 batch items per warp, each warp does 128 iterations. In each iteration, the 32 threads read 32 floats = 128 bytes, coalesced.
So kernel design:
```cuda
// Grid: (num_groups) = 256
// Block: 256 threads (8 warps)
// Each block handles one group
// Warp w handles batch items [w * 128, (w+1) * 128)
__global__ void fused_groupnorm_swish_mul_swish(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ gamma, // GroupNorm gamma (out_features)
const float* __restrict__ beta, // GroupNorm beta (out_features)
float* __restrict__ output,
int batch_size, int out_features, int num_groups, float eps) {
int group = blockIdx.x;
int num_channels = out_features / num_groups;
int group_start = group * num_channels;
int warps_per_block = blockDim.x / 32;
int warp_id = threadIdx.x / 32;
int lane_id = threadIdx.x % 32;
int channel_in_group = lane_id;
int batches_per_warp = batch_size / warps_per_block; // = 128
// Each warp computes partial sums for its batch items
float local_sum = 0.0f;
float local_sum_sq = 0.0f;
for (int b = 0; b < batches_per_warp; b++) {
int batch = warp_id * batches_per_warp + b;
if (batch < batch_size && channel_in_group < num_channels) {
int idx = batch * out_features + group_start + channel_in_group;
float val = input[idx];
local_sum += val;
local_sum_sq += val * val;
}
}
// Warp reduction
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
local_sum_sq += __shfl_down_sync(0xFFFFFFFF, local_sum_sq, offset);
}
// Now warp 0 has sums from all warps... need block reduction
__shared__ float shared_sum[8]; // for 8 warps
__shared__ float shared_sum_sq[8];
__shared__ float group_mean;
__shared__ float group_inv_std;
if (lane_id == 0) {
shared_sum[warp_id] = local_sum;
shared_sum_sq[warp_id] = local_sum_sq;
}
__syncthreads();
// One warp reduces the partial sums
if (warp_id == 0) {
float total_sum = (lane_id < warps_per_block) ? shared_sum[lane_id] : 0.0f;
float total_sq = (lane_id < warps_per_block) ? shared_sum_sq[lane_id] : 0.0f;
for (int offset = 16; offset > 0; offset /= 2) {
total_sum += __shfl_down_sync(0xFFFFFFFF, total_sum, offset);
total_sq += __shfl_down_sync(0xFFFFFFFF, total_sq, offset);
}
if (lane_id == 0) {
int N = batch_size * num_channels;
group_mean = total_sum / N;
float variance = total_sq / N - group_mean * group_mean;
group_inv_std = rsqrtf(variance + eps);
}
}
__syncthreads();
// Second pass: normalize and apply ops
for (int b = 0; b < batches_per_warp; b++) {
int batch = warp_id * batches_per_warp + b;
if (batch < batch_size && channel_in_group < num_channels) {
int idx = batch * out_features + group_start + channel_in_group;
int channel = group_start + channel_in_group;
float x = input[idx];
x = (x - group_mean) * group_inv_std;
// Apply gamma and beta (learned per-channel)
x = x * gamma[channel] + beta[channel];
// Swish 1: x * sigmoid(x)
float sig1 = 1.0f / (1.0f + expf(-x));
x = x * sig1;
// Multiply by weight
x = x * weight[channel];
// Swish 2
float sig2 = 1.0f / (1.0f + expf(-x));
x = x * sig2;
output[idx] = x;
}
}
}
```
Wait, I need to be more careful. In the reduction, each warp computes sum and sum_sq for 128 batch items * 1 channel. There are 32 threads in the warp, each handling a different channel. So after the warp shuffle reduction, lane_id 0 has the total sum and sum_sq for 128 batch items across 32 channels. Then we store to shared memory.
After all 8 warps store, warp 0 reads all 8 partial sums and reduces. Then lane_id 0 computes the group mean and inv_std.
Then in the second pass, each thread normalizes and applies the elementwise ops.
This looks correct. However, I realize there's a subtle issue: `channel_in_group = lane_id`, but lane_id goes from 0 to 31. We have 32 channels per group, so that's fine. But what if num_channels < 32? In our case num_channels = 32, so it's exactly 1 warp per "channel set".
Actually, this creates divergence if `num_channels < 32`, but for our case it's exactly 32.
Let me also think about the GroupNorm weight and bias. PyTorch's GroupNorm has optional weight and bias, but if not specified it uses learnable parameters... Wait, in the reference model:
```python
self.group_norm = nn.GroupNorm(num_groups, out_features)
```
In PyTorch, `nn.GroupNorm` by default has `affine=True`, so it has learnable per-channel weight and bias. I need to extract these from the module and pass them to the kernel.
Also, `eps` default is 1e-5.
Now, one concern: can we get even more performance by using vectorized loads/stores?
For 32 channels per group, each warp can load 32 floats for one batch item. That's exactly 128 bytes, which is a perfect L2 cache line! We can use `float4` loads to load 4 floats per thread (16 bytes), processing 8 batch items at a time or something.
Actually, using float4:
- 256 threads = 8 warps
- Instead of each warp processing 128 batch items × 1 channel, we can have each thread load a `float4`
- But the memory layout is: batch b has channels [g*32, g*32+31] contiguous
- If thread 0 loads floats at [0, 1, 2, 3] (batch b, channels 0,1,2,3), thread 1 loads [4, 5, 6, 7], etc.
- So 8 threads loading float4 loads 32 floats = 128 bytes.
With 256 threads, we could process 32 batch items in parallel (256 threads / 8 threads_per_batch = 32 batches), each loading a float4.
Actually let me simplify. With blockDim.x = 256:
- Process batch items in groups.
- For one "row" of 32 batch items:
- Batch b+0, channels [g*32, g*32+31]: 32 floats
- Batch b+1, channels [g*32, g*32+31]: 32 floats
- etc.
Hmm, vectorization might complicate the reduction. Let me first write a simple version and see if it compiles and is correct, then optimize further.
Wait, I realize I should think about whether we can avoid the explicit GEMM entirely and instead do the linear transform inside the kernel. But that's "fusing the GEMM" which I decided against.
Actually, let me reconsider. Since `torch.nn.Linear` is just `x @ W.T + b`, and x is (1024, 8192), W is (8192, 8192), b is (8192,):
- The output must be computed before GroupNorm
- We could call `torch.addmm(b, x, W.T)` or use the linear layer weight directly
In my `ModelNew.__init__`, I should copy the weights from the reference model. But wait, the evaluation framework creates `ModelNew(*get_init_inputs())` and then tests against `Model(*get_init_inputs())` with the same random seed? Actually, let me re-read...
The problem says "Your `ModelNew` class must: 1. Accept the same constructor arguments as `Model`." It doesn't say the framework will initialize it with the same weights. Let me think...
Usually in these benchmarks, the model is instantiated and weights are presumably randomly initialized, but then forward is tested. For correctness, we probably need the same weights. Let me check the reference: it creates `self.gemm = nn.Linear(...)`, `self.group_norm = nn.GroupNorm(...)`, etc. For correctness against reference, the framework probably:
1. Creates `Model` and `ModelNew` with same init args
2. Copies weights from `Model` to `ModelNew`, or initializes both with same seed
3. Runs forward on same input and compares
Actually, looking at the template more carefully, the evaluation framework likely initializes both models independently but with the same random seed, OR it might directly compare outputs of ModelNew against a pre-instantiated reference.
In any case, I should make sure my `ModelNew` has the same structure and learnable parameters as `Model`, or I should copy parameters in `__init__`.
Wait, actually the simplest approach: in `ModelNew.__init__`, create the same `nn.Linear`, `nn.GroupNorm`, and `nn.Parameter`. But then in `forward`, instead of using them as-is, extract their weights and call the custom kernel.
Actually, even simpler: we can keep the PyTorch modules but fuse the operations in forward. Let me do:
```python
class ModelNew(nn.Module):
def __init__(self, in_features, out_features, num_groups, multiply_weight_shape):
super().__init__()
self.gemm = nn.Linear(in_features, out_features)
self.group_norm = nn.GroupNorm(num_groups, out_features)
self.multiply_weight = nn.Parameter(torch.randn(multiply_weight_shape))
def forward(self, x):
# Use the weight/bias from self.gemm
x = torch.matmul(x, self.gemm.weight.t())
if self.gemm.bias is not None:
x = x + self.gemm.bias
# Call fused kernel
...
```
Wait, but I need to make sure weights match the reference. If the framework creates `Model` and `ModelNew` with the same random seed, they'll have identical weights. But if it creates `Model`, runs it, then creates `ModelNew` and expects the same outputs without copying weights... then the weights would differ due to different random initialization!
Actually, I think the standard eval setup is:
1. `ref = Model(*get_init_inputs())` with a fixed seed
2. `new = ModelNew(*get_init_inputs())` with the same fixed seed
3. Copy weights if necessary, or they already match
4. Run both with same input and compare
To be safe, I should make `ModelNew` instantiate the same submodules so random initialization is identical when using the same seed.
But then in `forward`, I don't need to use `self.gemm` as an `nn.Linear` - I can do the matmul manually and then call the fused kernel for the rest.
For the fused kernel, I need to pass:
- input (GEMM output)
- GroupNorm weight and bias `(out_features,)`
- multiply_weight `(out_features,)`
- output buffer
- batch_size, out_features, num_groups, eps
Now let me think about register usage and occupancy. With 256 threads per block and simple math, we should get good occupancy.
Let me also think about whether I should use the inline `load_inline` and how to structure it.
Another idea: I realized that on H100, we might get much better performance by using `torch.compile` with fullgraph... but the point of the exercise is to write a custom CUDA kernel.
Also, since the GroupNorm module's weight and bias are per-channel, and multiply_weight is also per-channel, we can combine some of these operations.
Actually: `normalized * gamma + beta` then `* sigmoid(...)` then `* weight` then `* sigmoid(...)`
Can we precompute `gamma` and `weight` together? No, because the sigmoid is non-linear and depends on the intermediate values.
Let me write the CUDA code now.
One more thing: I should think about whether we can reduce warp divergence. The kernel has `if (batch < batch_size && channel_in_group < num_channels)`. For our case, batch_size=1024, num_channels=32. With 8 warps and 128 batches per warp, and 32 channels per group (exactly 1 warp), all threads in all warps are active since 8*128=1024 exactly. Actually wait, `b` goes from 0 to `batches_per_warp - 1 = 127`, and `batch = warp_id * 128 + b`. For warp_id 0..7, batch = 0..127, 128..255, ..., 896..1023. That's exactly 1024 batch items. But `channel_in_group = lane_id` goes from 0..31, and `num_channels = 32`, so all threads pass the condition. No divergence!
Now, for `__shfl_down_sync`, since we're reducing within a warp and all 32 threads in the warp are active, shuffles work perfectly.
But wait, after the warp reduction using shuffle, `lane_id 0` has the correct sum. Then we store to shared memory. What about other lanes? They have who-knows-what, but we only store from lane_id == 0, which is fine.
Hmm, actually there's an important issue. In the first pass, each thread in the warp computes sum for ONE channel. After warp shuffle, lane 0 has the sum of all 32 channels for 128 batches. That's:
`total_sum_for_warp = sum over all 128 batches and all 32 channels in the group`
And there are 8 warps, each covering 128 batches. So the block-level reduction gives us the sum over all 1024 batches and all 32 channels. That's correct!
Wait, but each thread in the warp is reading 128 values (one channel across 128 batches). There are 32 threads, so total reads per warp = 32 * 128 = 4096 values. 8 warps = 32768 values total. That's exactly the group size. Good.
Now for the second pass, same pattern. Each thread reads 128 values. After normalization, it writes back 128 values. Total 32768 reads and 32768 writes per block.
Memory traffic per block: 32768 * 4 bytes read + 32768 * 4 bytes write = 262 KB read + 262 KB write = 524 KB.
For 256 blocks: 134 MB read + 134 MB write = 268 MB total.
At 3.35 TB/s: ~80 microseconds.
Plus the GEMM. A 1024x8192x8192 GEMM at ~67 TFLOPS: 1024*8192*8192*2 / 67e12 = 2.03 ms. Actually wait, cuBLAS on H100 with Tensor Cores should be much faster, even for FP32. H100 has 989 TFLOPS with TF32 Tensor Cores. At that rate: 1024*8192*8192*2 / 989e12 = 0.138 ms!
Actually, the reference model would use cuBLAS which uses Tensor Cores when available. So the GEMM might be around 0.2-0.5ms.
Then the unfused PyTorch ops:
- GroupNorm: reads 32MB, writes 32MB, plus reduction ~20 microseconds
- Swish: reads 32MB, writes 32MB ~10 microseconds
- Mul: reads ~32MB, writes 32MB ~10 microseconds
- Swish: reads 32MB, writes 32MB ~10 microseconds
Total for unfused elementwise: maybe 50-100 microseconds.
Fused: maybe 30-50 microseconds.
So our win might be small unless PyTorch's GroupNorm is particularly slow. PyTorch GroupNorm can sometimes launch multiple kernels and not be fully optimized.
Actually, in PyTorch, GroupNorm for 2D input might be implemented using multiple kernels or less than optimal reduction. Our fused kernel should still win.
But here's a bigger optimization: what if we can use FP16 Tensor Cores for the GEMM and do the rest in FP32? The problem says "numerically equivalent outputs at fp32 precision" with tolerance 1e-4. But the GEMM result from FP16 vs FP32 might differ by more than 1e-4 relative for large matrices... Actually, for a (1024, 8192) @ (8192, 8192) GEMM, using FP16 Tensor Cores and accumulating in FP32, the absolute error could be tens or hundreds. That's way more than 1e-4. So we can't use FP16 for the GEMM.
We could use TF32, which has about 1e-3 relative precision. Again, for large GEMM with accumulation, the error might exceed 1e-4.
So let's keep GEMM in FP32. PyTorch's Linear on H100 will use TF32 by default though! Let me check... PyTorch's behavior is:
- `torch.backends.cuda.matmul.allow_tf32` defaults to True on Ampere/Hopper for FP32.
- If True, FP32 inputs use TF32 Tensor Cores.
So the reference model actually might be using TF32! In that case, we can also use `torch.matmul` with TF32 and it should match exactly. Actually, `nn.Linear` uses the same backend.
And actually, for correctness, our `ModelNew` must match the `Model` output. So if `Model` uses TF32 for the GEMM (due to PyTorch defaults), we should also use TF32 for the GEMM. Using `torch.matmul(x, self.gemm.weight.t())` should automatically use whatever settings PyTorch has.
Let me write the kernel code now. One more optimization: can we process more than one channel per thread? Since 32 threads exactly cover 32 channels, using float4 would let each thread process 4 channels, so we'd need only 8 threads for the channels. Let's explore vectorization.
Actually, with float4:
- num_channels = 32
- Each float4 covers 4 channels
- We need 32/4 = 8 "channel slots" per batch
- With 256 threads, we can cover 256/8 = 32 batch items simultaneously
- 1024 / 32 = 32 iterations
This might be faster due to vectorized memory access. Let me think about how to organize it.
Memory layout: For batch b, group g, channels [g*32, g*32+31] are contiguous.
- Thread i (0..255): processes batch b = i / 8, channel_offset = (i % 8) * 4
- Thread i loads `float4` from `input[b * out_features + group_start + channel_offset]`
- This means threads in the same warp load different batches, i.e., warp 0 loads 32 batches (since 256 threads / 8 = 32 batches per "full set"... wait)
Actually with 256 threads and 8 threads per batch (for 32 channels with float4):
- Threads 0..7: batch 0, channels 0..3, 4..7, ..., 28..31
- Threads 8..15: batch 1, channels 0..3, 4..7, ..., 28..31
- ...
- Threads 248..255: batch 31, channels 0..3, 4..7, ..., 28..31
But these threads are all in the same block. Warp 0 has threads 0..31, which covers batches 0..3 (32 threads / 8 = 4 batches per warp).
- Within warp 0:
- threads 0..7: batch 0
- threads 8..15: batch 1
- threads 16..23: batch 2
- threads 24..31: batch 3
For coalescing: threads 0..7 load batch 0, channels 0..3, 4..7, ..., 28..31. The addresses are:
- thread 0: batch0, channels 0..3 → address = b*8192 + g*32 + 0
- thread 1: batch0, channels 4..7 → address = b*8192 + g*32 + 4
- ...
- thread 7: batch0, channels 28..31 → address = b*8192 + g*32 + 28
These are 16 bytes apart each, spanning 128 bytes. But are they coalesced? For H100, coalescing requires threads in a warp to access contiguous memory. Here, threads 0..7 access [0..3], [4..7], ..., [28..31] for the SAME batch. Threads 8..15 access the same channels but for batch 1, which is at b*8192 + ... which is far away.
So within warp 0:
- threads 0..7 access batch 0, contiguous 128B
- threads 8..15 access batch 1, contiguous 128B but separated by 8192*4 = 32KB
This is NOT coalesced within the warp! A warp accesses two different 128B regions for two different batches. The GPU will issue two memory transactions (one for batch 0, one for batch 1), each 128B. This is still fine (no wasted bandwidth), but it's not a single transaction.
Actually, for better coalescing, we want all threads in a warp to access the same 128B cache line if possible. But here, the data for batch 0 and batch 1 are in different cache lines (32KB apart), so they can't be in the same line anyway.
So the float4 approach is memory-bandwidth-optimal even if not "coalesced" in the strictest sense. The GPU will issue separate transactions for each 128B chunk.
Let me go back to the simpler scalar approach first, which has perfect coalescing. The scalar kernel with 32 threads per warp accessing the same batch's 32 channels contiguously is perfectly coalesced.
With 256 threads = 8 warps, each warp handles 128 batches:
- Warp 0: batch 0..127, each warp iteration processes one batch
- In iteration b, threads 0..31 read channels 0..31 of batch b
- These addresses differ by 4 bytes → 128-byte contiguous region → 1 transaction
This is perfectly coalesced and simple. Let me use this.
Actually, wait. I said "8 warps = 256 threads". But I also said "each warp handles 128 batches". With 32 threads per warp and each thread reading 128 items (one per batch), that's 32 * 128 = 4096 reads per warp. 8 warps = 32768 reads total. Good.
And only the warp-level reduction needs shuffle, then block-level in shared memory.
One more consideration: occupancy. 256 threads per block is good and gives 2 blocks per SM on H100 (max 2048 threads per SM). Actually H100 has 2048 threads/SM, so 256 threads/block = 8 possible blocks/SM. But limited by registers too. With our kernel using maybe 32 registers, we can probably get 4+ blocks/SM.
Let me write the Python code.
Wait, I need to make sure I understand how to use `torch.utils.cpp_extension.load_inline` correctly.
```python
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
__global__ void fused_kernel(...) {
...
}
torch::Tensor fused_forward(
torch::Tensor input,
torch::Tensor gamma,
torch::Tensor beta,
torch::Tensor weight,
float eps) {
int batch_size = input.size(0);
int out_features = input.size(1);
int num_groups = gamma.size(0); // wait, gamma is (out_features,) not (num_groups,)
...
}
"""
module = load_inline(
name="fused_ops",
cuda_sources=cuda_source,
functions=["fused_forward"],
extra_cuda_cflags=["-O3", "--use_fast_math"]
)
```
Wait, gamma and beta for GroupNorm are per-channel, shape `(out_features,)`. The number of groups is passed separately.
Let me also think about whether I should handle arbitrary num_groups, or hardcode for this specific case. For robustness, I'll handle arbitrary values but the main optimization is for num_channels % 32 == 0 or similar.
Actually, for simplicity and maximum performance on this specific problem, I'll write code optimized for the exact dimensions: batch_size=1024, out_features=8192, num_groups=256, channels_per_group=32. But I'll keep it general enough to work for other sizes.
Let me code this up now.
Wait, I realize there's another optimization. I can precompute the gamma and multiply_weight products, but since there's a non-linear activation in between (sigmoid), I can't merge them. However, I notice the pattern is:
```
x = groupnorm(x)
x = x * sigmoid(x)
x = x * weight
x = x * sigmoid(x)
```
Can I combine the two sigmoid evaluations? Not easily since the second sigmoid depends on the first Swish and multiply.
Another thought: instead of doing `1.0f / (1.0f + expf(-x))`, I can use `sigmoid(x)` from CUDA which might have an intrinsic. Actually, `1.0f / (1.0f + __expf(-x))` with `--use_fast_math` is fast.
Actually, there's `__fmaf_rn` and similar, but for now let's just use standard arithmetic with `--use_fast_math`.
One more thought: instead of reading `input` twice (once for reduction, once for normalization), could I store the data in shared memory between the two phases? 32768 floats = 131KB. H100 SM has 228KB shared memory. That's enough to fit all data for one group in shared memory!
If I load all data into shared memory in the first pass, then:
1. First pass: load from global memory into shared memory, simultaneously compute partial sums
2. Compute mean/inv_std
3. Second pass: read from shared memory, normalize, apply ops, write to global memory
This saves one read from global memory! But wait, we need all data in shared memory before we can normalize, and we're computing partial sums while loading. That works!
Actually, we can load in chunks. Each warp loads its portion, stores to smem, and accumulates partial sum. Then after all data is in smem, do the reduction, compute stats, then read from smem to apply ops.
32768 floats * 4 bytes = 131KB. Under the 227KB limit. Perfect!
But actually, with 256 threads and 131KB smem, and up to 32 blocks per SM (but limited by smem to 1 block per SM since 131KB * 2 > 227KB), we get at most 1 block per SM. That means 132 SMs on H100 = 132 blocks in parallel. We have 256 blocks total, so it would take 2 waves. That's still okay.
Wait, H100 SXM5 has 132 SMs, but some might be disabled. Even with 1 block per SM, 256 blocks take about 2 waves.
But is 1 block per SM a problem? 256 threads per block = 8 warps, which is decent utilization. But if we use less smem, we can fit more blocks per SM.
Alternative: don't keep all data in smem. Compute stats with two passes over global memory:
- Pass 1: compute mean and inv_std
- Pass 2: read again, normalize, write
Global memory bandwidth is 3.35 TB/s. Reading 32MB twice is 64MB = ~20 microseconds. Using smem to avoid the second read saves ~10 microseconds. The tradeoff is lower occupancy (1 block per SM vs maybe 4-8 blocks per SM).
Actually, with 131KB smem per block and 227KB per SM, we can only fit 1 block per SM. With 256 threads/block and 1 block/SM = 256 threads/SM. H100 can run 2048 threads/SM, so we're at 12.5% thread occupancy! That's bad for latency hiding.
Without smem caching, we might use ~4KB smem (just for the reduction). Then we can fit 32 blocks per SM (limited by 32 blocks/SM max), or 2048 threads/SM / 256 = 8 blocks per SM. So 8x higher occupancy!
At 3.35 TB/s, reading 32MB is ~10 microseconds. Reading it twice is ~20 microseconds. With 8x occupancy but 2x more reads, the global memory approach might actually be better because we can hide latency with more warps across more blocks.
Let me think about this more carefully. On H100:
- 132 SMs
- Each SM can have up to 32 blocks, 2048 threads
- Achieved occupancy depends on how many warps can run concurrently
If we have 8 blocks per SM (256 threads each = 2048 threads), that's full thread capacity. BUT with shared memory, we'd be limited to 1 block per SM, so only 256 threads. That's 8x fewer active warps for memory latency hiding.
Given that global memory latency is ~400-600 cycles on H100, and at 2 GHz that's ~200-300ns, and bandwidth is very high, I think the lower-SMEM approach with higher occupancy is better.
Actually, let me reconsider. Even with 1 block per SM and 256 threads, we have 8 warps. At H100's memory throughput, 32MB reads + 32MB writes = 64MB total. At 3.35 TB/s = 64 / 3350 = 0.019 ms = 19 microseconds. Can 8 warps hide 19 microseconds of memory latency?
Each memory instruction has ~300ns latency. With 8 warps and assuming each warp issues memory requests quickly, we might have ~8 * 32 = 256 threads worth of memory requests in flight. That's probably not enough to fully saturate the memory system, but it might be close.
With 8 blocks per SM (2048 threads = 64 warps), we have much more memory-level parallelism. The extra read pass costs 10 microseconds but with 8x more warps the memory system is better utilized.
I think I'll start without the smem cache and just do two passes over global memory.
Wait, actually there's another idea. With batch_size=1024 and channels_per_group=32, and 8 warps per block, each warp processes 128 batches. We can use a persistent kernel or process multiple groups per block. For instance, with 256 threads, we could process 2 groups (128 threads per group, or 4 warps per group). Then each block uses 2x fewer blocks. But then we'd have 128 blocks instead of 256, and could potentially have 2 blocks per SM, giving 512 threads per SM. That's still not full but better.
With 4 warps per group: 4 warps * 32 threads = 128 threads per group. Each warp handles 256 batches. That still works.
Let me stick with 1 group per block for simplicity.
You know what, let me also consider another approach: each thread block handles a tile of (batch, channels), and we can use a grid-stride loop for generality.
Actually, let me first write a basic version, compile it, check correctness, then optimize.
One technical detail: I need to extract gamma and bias from `self.group_norm`. In PyTorch, `nn.GroupNorm` stores them as `self.group_norm.weight` and `self.group_norm.bias`.
And for the linear layer, I'll do:
```python
x = torch.matmul(x, self.gemm.weight.t())
if self.gemm.bias is not None:
x = x + self.gemm.bias
```
This should give the same result as `self.gemm(x)`.
Actually, wait. `nn.Linear` applies the bias with broadcasting. `torch.matmul(x, weight.t()) + bias` does the same.
Now for the kernel interface. Since `load_inline` auto-generates the pybind11 module, I just need to define a C++ function that takes torch Tensors and returns a torch Tensor.
Let me also think about whether I need to handle non-contiguous tensors. For safety, I'll check or use `.contiguous()`.
Let me draft the code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void fused_gn_swish_mul_swish_kernel(
const float* __restrict__ input,
const float* __restrict__ gamma,
const float* __restrict__ beta,
const float* __restrict__ mul_weight,
float* __restrict__ output,
int batch_size,
int out_features,
int num_groups,
float eps) {
int group = blockIdx.x;
int tid = threadIdx.x;
int warp_id = tid / 32;
int lane_id = tid % 32;
int num_warps = blockDim.x / 32;
int channels_per_group = out_features / num_groups;
int group_start = group * channels_per_group;
int batches_per_warp = (batch_size + num_warps - 1) / num_warps;
int batch_start = warp_id * batches_per_warp;
int batch_end = min(batch_start + batches_per_warp, batch_size);
// Handle case where channels_per_group < 32
if (lane_id >= channels_per_group) return;
float local_sum = 0.0f;
float local_sum_sq = 0.0f;
// First pass: compute partial sums
for (int batch = batch_start; batch < batch_end; batch++) {
int idx = batch * out_features + group_start + lane_id;
float val = input[idx];
local_sum += val;
local_sum_sq += val * val;
}
// Warp-level reduction
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
local_sum_sq += __shfl_down_sync(0xFFFFFFFF, local_sum_sq, offset);
}
// Block-level reduction using shared memory
extern __shared__ float smem[];
float* s_sum = smem;
float* s_sum_sq = smem + num_warps;
if (lane_id == 0) {
s_sum[warp_id] = local_sum;
s_sum_sq[warp_id] = local_sum_sq;
}
__syncthreads();
float mean, inv_std;
if (warp_id == 0) {
float total_sum = (lane_id < num_warps) ? s_sum[lane_id] : 0.0f;
float total_sq = (lane_id < num_warps) ? s_sum_sq[lane_id] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
total_sum += __shfl_down_sync(0xFFFFFFFF, total_sum, offset);
total_sq += __shfl_down_sync(0xFFFFFFFF, total_sq, offset);
}
if (lane_id == 0) {
int N = batch_size * channels_per_group;
mean = total_sum / N;
float variance = total_sq / N - mean * mean;
inv_std = rsqrtf(variance + eps);
s_sum[0] = mean;
s_sum[1] = inv_std;
}
}
__syncthreads();
mean = s_sum[0];
inv_std = s_sum[1];
// Second pass: normalize, swish, multiply, swish
for (int batch = batch_start; batch < batch_end; batch++) {
int idx = batch * out_features + group_start + lane_id;
int channel = group_start + lane_id;
float x = input[idx];
x = (x - mean) * inv_std;
x = x * gamma[channel] + beta[channel];
float sig1 = 1.0f / (1.0f + expf(-x));
x = x * sig1;
x = x * mul_weight[channel];
float sig2 = 1.0f / (1.0f + expf(-x));
x = x * sig2;
output[idx] = x;
}
}
torch::Tensor fused_forward(
torch::Tensor input,
torch::Tensor gamma,
torch::Tensor beta,
torch::Tensor mul_weight,
float eps) {
auto batch_size = input.size(0);
auto out_features = input.size(1);
auto num_groups = gamma.size(0) / (out_features / gamma.size(0)); // hmm, need num_groups somehow
// Actually gamma is (out_features,), we need num_groups passed in
}
```
Wait, I need to pass `num_groups` to the kernel. Let me adjust the function signature.
Actually, I can just pass it as an argument from Python.
Let me fix the signature:
```cpp
torch::Tensor fused_forward(
torch::Tensor input,
torch::Tensor gamma,
torch::Tensor beta,
torch::Tensor mul_weight,
int num_groups,
float eps)
```
And in Python:
```python
output = fused_forward(gemm_out, self.group_norm.weight, self.group_norm.bias, self.multiply_weight, self.num_groups, self.group_norm.eps)
```
Actually, `self.group_norm.eps` defaults to 1e-5 in PyTorch, but let me pass it explicitly.
Wait, there's a problem. I declared `extern __shared__ float smem[];` for dynamic shared memory. But I need to calculate the size. For `num_warps = blockDim.x / 32`, the size is `2 * num_warps * sizeof(float)` for s_sum and s_sum_sq. Then I reuse the first few entries for mean and inv_std.
In the kernel launch, I'd do something like:
```cpp
fused_kernel<<<num_groups, threads_per_block, 2 * num_warps * sizeof(float)>>>(...)
```
And `threads_per_block` needs to be at least 32 and ideally a multiple of 32. For 1024 batch items and 32 channels, we'd want enough warps to cover all batch items. If we use 1024 threads (32 warps), each warp handles 32 batches. That's 32 batches * 32 channels = 1024 elements per warp, total 32768. But 1024 threads might be overkill and reduce occupancy.
Let's use 256 threads (8 warps), each warp handles 128 batches.
Actually, 1024 threads would mean 32768 / 1024 = 32 elements per thread, which is fine. But 1024 threads per block means max 2 blocks per SM (2048 threads/SM), giving 512 blocks active across 256 SMs... wait, we only have 256 blocks total. So it doesn't matter much.
Let me use 256 threads for good balance.
One more issue: I used `__shfl_down_sync` but only for lane_id < num_warps in the block reduction. With `num_warps = 8`, lanes 0..7 participate, lanes 8..31 just passively participate in the shuffle but won't contribute because they read 0. The `__activemask()` won't be needed since all threads in warp 0 are active and the default mask 0xFFFFFFFF is fine.
Actually, for correctness with `__shfl_down_sync`, all threads specified in the mask must execute the instruction. In our case, warp 0 has all 32 threads active and they all execute the shuffle. Threads with lane_id >= num_warps just have total_sum = 0, which is fine.
But for the first shuffle, the condition `if (lane_id < num_warps) s_sum[lane_id] else 0.0f` ensures non-participating lanes have 0.
One thing: I need to make sure `num_warps` is a compile-time constant or the loop unrolling works. Since `num_warps = blockDim.x / 32` at runtime, the `#pragma unroll` on the shuffle loop with variable limit won't fully unroll... wait, offset goes 16, 8, 4, 2, 1 regardless of num_warps. But if num_warps < the number of iterations needed, the extra iterations still work but might be wasted. Let's leave it as is.
Wait, the shuffle loop always does 16, 8, 4, 2, 1, which is 5 iterations. That's only 5 instructions. Fine.
Okay, let me also handle the case where `batch_size` is not divisible by `num_warps`. I already do `batch_end = min(batch_start + batches_per_warp, batch_size)`. With `batches_per_warp = ceil(batch_size / num_warps)`, some warps might not have work, but since `batch_end` caps it, threads in those warps won't enter the loop. That's okay.
But actually, there's divergence: if `batch_start >= batch_size`, the loop doesn't execute and `local_sum = 0`. The shuffle still works fine. No divergence within warps since all threads in a warp have the same `batch_start` and `batch_end` (they have the same `warp_id`). Good.
Now let me check: with blockDim.x = 256, num_warps = 8. batch_size = 1024. batches_per_warp = 128. Each warp processes exactly 128 batches. No divergence.
channels_per_group = 32. lane_id from 0..31. If lane_id >= 32... wait, lane_id is always 0..31. Since channels_per_group = 32, the `if (lane_id >= channels_per_group) return;` doesn't trigger.
Great. Now I need to write the full Python module.
One more optimization: I can use `__expf` instead of `expf` with `--use_fast_math`. Actually, `--use_fast_math` already maps `expf` to `__expf`, so no need to change.
Also, I should check if I can use `fast_math` without losing too much precision. For fp32 with tolerance 1e-4, fast math should be fine.
Let me also consider whether I should vectorize the stores. For 32 contiguous floats, a warp can write them with each thread writing 1 float. That's already a single 128-byte transaction (or maybe 2x 64-byte, depending on architecture). On H100, the L2 cache line is 128 bytes, so 32 threads writing 32 consecutive floats = 128 bytes = 1 cache line. Perfect.
Can we do vectorized loads/stores with `float4`? Each thread would read 4 floats. 8 threads per batch × 4 floats = 32 floats. But then we'd need only 8 threads for channels, and the other 24 threads in the warp would be idle. That's bad warp utilization.
Actually, we could have all 32 threads in a warp each read `float4`, covering 4 batches of 32 floats each. But the batches are 32KB apart in memory, so this is 4 separate cache lines. Still fine for bandwidth. But the processing then becomes more complex.
I think scalar access with good coalescing is the way to go.
Now let me write the full code. I also need to consider whether I should use `torch::empty` or `torch::zeros` for the output. `torch::empty` is fine.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void fused_kernel(
const float* __restrict__ input,
const float* __restrict__ gamma,
const float* __restrict__ beta,
const float* __restrict__ mul_weight,
float* __restrict__ output,
int batch_size,
int out_features,
int num_groups,
float eps) {
int group = blockIdx.x;
int channels_per_group = out_features / num_groups;
int group_start = group * channels_per_group;
int tid = threadIdx.x;
int warp_id = tid >> 5;
int lane_id = tid & 31;
int num_warps = blockDim.x >> 5;
if (lane_id >= channels_per_group) return;
int batches_per_warp = (batch_size + num_warps - 1) / num_warps;
int batch_begin = warp_id * batches_per_warp;
int batch_end = min(batch_begin + batches_per_warp, batch_size);
float local_sum = 0.0f;
float local_sum_sq = 0.0f;
#pragma unroll 4
for (int batch = batch_begin; batch < batch_end; ++batch) {
int idx = batch * out_features + group_start + lane_id;
float val = input[idx];
local_sum += val;
local_sum_sq += val * val;
}
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
local_sum_sq += __shfl_down_sync(0xFFFFFFFF, local_sum_sq, offset);
}
extern __shared__ float s_mem[];
float* s_sum = s_mem;
float* s_sum_sq = s_mem + num_warps;
if (lane_id == 0) {
s_sum[warp_id] = local_sum;
s_sum_sq[warp_id] = local_sum_sq;
}
__syncthreads();
if (warp_id == 0) {
float total_sum = (lane_id < num_warps) ? s_sum[lane_id] : 0.0f;
float total_sq = (lane_id < num_warps) ? s_sum_sq[lane_id] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
total_sum += __shfl_down_sync(0xFFFFFFFF, total_sum, offset);
total_sq += __shfl_down_sync(0xFFFFFFFF, total_sq, offset);
}
if (lane_id == 0) {
int N = batch_size * channels_per_group;
float mean = total_sum / N;
float variance = total_sq / N - mean * mean;
s_sum[0] = mean;
s_sum[1] = rsqrtf(variance + eps);
}
}
__syncthreads();
float mean = s_sum[0];
float inv_std = s_sum[1];
#pragma unroll 4
for (int batch = batch_begin; batch < batch_end; ++batch) {
int idx = batch * out_features + group_start + lane_id;
int channel = group_start + lane_id;
float x = input[idx];
x = (x - mean) * inv_std;
x = x * gamma[channel] + beta[channel];
float sig = 1.0f / (1.0f + expf(-x));
x = x * sig;
x = x * mul_weight[channel];
sig = 1.0f / (1.0f + expf(-x));
x = x * sig;
output[idx] = x;
}
}
torch::Tensor fused_forward(
torch::Tensor input,
torch::Tensor gamma,
torch::Tensor beta,
torch::Tensor mul_weight,
int num_groups,
float eps) {
auto batch_size = input.size(0);
auto out_features = input.size(1);
auto output = torch::empty_like(input);
int threads_per_block = 256;
int num_warps = threads_per_block / 32;
fused_kernel<<<num_groups, threads_per_block, num_warps * 2 * sizeof(float)>>>(
input.data_ptr<float>(),
gamma.data_ptr<float>(),
beta.data_ptr<float>(),
mul_weight.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
out_features,
num_groups,
eps);
return output;
}
"""
fused_module = load_inline(
name="fused_gn_swish",
cuda_sources=cuda_source,
functions=["fused_forward"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
)
class ModelNew(nn.Module):
def __init__(self, in_features, out_features, num_groups, multiply_weight_shape):
super().__init__()
self.gemm = nn.Linear(in_features, out_features)
self.group_norm = nn.GroupNorm(num_groups, out_features)
self.multiply_weight = nn.Parameter(torch.randn(multiply_weight_shape))
self.num_groups = num_groups
def forward(self, x):
x = torch.matmul(x, self.gemm.weight.t())
if self.gemm.bias is not None:
x = x + self.gemm.bias
x = fused_module.fused_forward(
x,
self.group_norm.weight,
self.group_norm.bias,
self.multiply_weight,
self.num_groups,
self.group_norm.eps
)
return x
```
Wait, I need to handle the case where `self.group_norm` has no weight/bias. In PyTorch, `nn.GroupNorm` defaults to `affine=True`, so weight and bias exist. But what if `affine=False`? The reference uses default `affine=True`, so we're fine.
But there's a subtle issue: `self.group_norm.eps` is a Python float, but in the reference it's stored as a float. That should be fine.
Let me also check: does `nn.GroupNorm` have an `eps` attribute? Yes, it's stored in `self.eps`.
Now, one concern. The PyTorch `nn.Linear` might add the bias inside its own cuBLAS call or via a separate kernel. When I do `torch.matmul(x, weight.t()) + bias`, the bias addition might be fused by the PyTorch JIT/cuBLAS or it might be a separate kernel. Either way, it's just one elementwise pass, and we'll still save time on the GroupNorm + Swish + Mul + Swish fusion.
But actually, can I use `torch.addmm` to fuse the bias addition with the GEMM?
```python
x = torch.addmm(self.gemm.bias, x, self.gemm.weight.t())
```
Yes, and this is exactly what `nn.Linear.forward` does under the hood! So using `torch.addmm` should give us the same behavior but explicitly.
Let me modify the forward to:
```python
x = torch.addmm(self.gemm.bias, x, self.gemm.weight.t())
```
Actually, I should double-check the argument order for `addmm`:
`torch.addmm(input, mat1, mat2, *, beta=1, alpha=1, out=None)` computes `beta * input + alpha * (mat1 @ mat2)`
So `torch.addmm(self.gemm.bias, x, self.gemm.weight.t())` gives `self.gemm.bias + x @ self.gemm.weight.t()`. Correct!
Now, one more thing. I should add `.contiguous()` to the inputs just in case, or I can assume they're contiguous since they come from `addmm` which outputs contiguous tensors.
Let me also make sure the types are right. Since the reference uses float32, I'll use float32.
Okay, let me compile and test this. I realize I should add some input validation or handle cases where inputs are not on CUDA? The reference is presumably on CUDA though. Let me check... the benchmark setup usually creates CUDA tensors. But to be safe, I won't add explicit device checks.
Actually, wait. I need to consider whether the evaluation might test with non-contiguous tensors. But since our inputs come from `get_inputs()` which does `torch.rand(...)`, they're contiguous. And our output is contiguous.
Let me compile this first.