reasoning
Let me analyze this problem:
1. The model does: Linear (GEMM) → GroupNorm → Min along dim=1, keepdim=True → Add bias
2. Input: [batch_size=1024, in_features=8192]
3. Output of Linear: [1024, 8192]
4. GroupNorm with 512 groups on 8192 channels: [1024, 8192]
5. Min along dim=1: [1024, 1] (takes min across channels for each batch item)
6. Add bias with shape (1, 8192, 1, 1): [1024, 1] + [1, 8192, 1, 1] = [1024, 8192, 1, 1]
Wait, let me re-read. After min, x has shape [1024, 1]. Then adding bias with shape (1, out_features, 1, 1) = (1, 8192, 1, 1).
In PyTorch broadcasting:
[1024, 1] + [1, 8192, 1, 1] → [1024, 8192, 1, 1]
Actually wait, let me trace through more carefully:
- x starts as [1024, 8192]
- gemm(x): [1024, 8192] @ [8192, 8192] + [8192] = [1024, 8192]
- group_norm(x): operates on [1024, 8192], treating as N=1024, C=8192. Since GN expects N,C,* spatial, it can work with [1024, 8192] as N=1024, C=8192, no spatial dims (or H=W=1). Output: [1024, 8192]
- torch.min(x, dim=1, keepdim=True)[0]: [1024, 1]
- x + self.bias: self.bias is [1, 8192, 1, 1], x is [1024, 1]
In broadcasting: [1024, 1] + [1, 8192, 1, 1]
- Expand [1024, 1] to [1024, 1, 1, 1]
- Expand [1, 8192, 1, 1] to [1024, 8192, 1, 1]
- Result: [1024, 8192, 1, 1]
So the final output shape is [1024, 8192, 1, 1].
This gives us optimization opportunities:
1. The min operation reduces each row to a single value
2. The result is then broadcast-added across all channels
3. GroupNorm is applied, then immediately followed by a min and broadcast add
Wait, but GroupNorm changes values per channel, then we take min across channels. This means we can't really fuse the GroupNorm with the min in a trivial way without computing the GN first.
However, we can still fuse the operations. Let me think about what each operation does:
GEMM: Y = X @ W^T + b (linear)
GroupNorm: For each sample, divide into 512 groups of 16 channels each. Each group has mean and std computed, then normalize. So channels 0-15 in group 0, 16-31 in group 1, etc. Wait, out_features=8192, num_groups=512, so 8192/512 = 16 channels per group.
Each group is normalized independently. Given affine parameters gamma and beta (learnable per channel).
Actually, let me reconsider. Can we avoid computing the full GroupNorm output?
After GroupNorm, we compute min across all channels for each batch item, then add bias and broadcast. So the final output for batch item i, channel j is:
output[i, j] = min_k(GN(x)[i, k]) + bias[j]
Hmm, that's interesting. Since we only need the min across channels after GN, and then add the same min (plus bias per channel) to each output channel... Wait no, we add bias[j] to each output position, where bias is per channel. But the min is the same for all channels in a given batch item.
So output[i, j, 0, 0] = min_val[i] + bias[0, j, 0, 0]
Where min_val[i] = min over k of GN(x)[i, k].
So the whole computation for batch i essentially produces a scalar min_val[i], then min_val[i] + bias[j] for each j.
This means we don't need to materialize all of GN(x)[i, k] and then take min - we can compute GN(x)[i, k] on the fly and track the minimum! But wait, GN involves computing mean and variance per group, which requires looking at all channels in the group first.
For GroupNorm with 512 groups and 16 channels per group:
- For each batch item and each group: compute mean of 16 values
- Compute variance of 16 values
- Normalize each of the 16 values: (x - mean) / sqrt(var + eps)
- Apply gamma and beta per channel
To get GN(x)[i, k], we need to know which group channel k belongs to, and have computed the mean/variance for that group.
Strategy:
1. For each batch item, compute the GEMM output (or we can fuse with the rest)
2. Compute group means and variances
3. Compute normalized values and track minimum per batch item
4. Compute final output = min_val + bias
Actually, step 1 might be the bottleneck. Let's consider sizes:
- GEMM: [1024, 8192] @ [8192, 8192] = 1024 * 8192 * 8192 * 2 ≈ 137B FLOPs
- This is quite large and compute-bound
GroupNorm on [1024, 8192]: mostly memory bandwidth bound, negligible compute
Min: trivial
Bias add: trivial
So the GEMM will dominate. On H100, 137B FLOPs at 67 TFLOPS = ~2ms. But we can use cuBLAS for this.
Actually, at 1024 * 8192 * 8192 * 2 = 137,438,953,472 FLOPs = 137 GFLOPs. H100 does 67 TFLOPS for fp32, so ~2ms for naive. With Tensor Cores in TF32, ~989 TFLOPS, but for precision we need fp32.
Let me think about this differently. The batch is 1024, which is small. Can we fit everything in shared memory or registers?
Wait, we need to think if we can fuse the GEMM with the reduction. For small M (1024) and large K,N (8192 each), the GEMM is M=1024, N=8192, K=8192. This is a classic square-ish GEMM. We should just call cuBLAS for this.
But then after the GEMM, we have [1024, 8192] which needs GroupNorm, min, and bias add. We can definitely fuse these three operations into a single kernel.
Actually, we might be able to use CUTLASS or custom GEMM + epilogue fusion. But let's start simpler: call cuBLAS for GEMM, then a custom kernel for GN + min + bias.
Wait, but the reference does:
1. Linear (cuBLAS GEMM)
2. GroupNorm (optimized PyTorch kernel)
3. Min (reduction)
4. Bias add (broadcasted elementwise)
PyTorch can already fuse some things. But our custom kernel can fuse steps 2-4, and even step 1 if we're clever.
Actually for H100, let me think about what would be fastest:
- cuBLAS GEMM for the linear layer
- Custom kernel that reads GEMM output, computes group norm on the fly, finds min, and writes the final broadcasted result
Since the output after min is [1024, 1] broadcast to [1024, 8192, 1, 1], but wait - the final shape has spatial dims 1x1. So effectively we write [1024, 8192].
For the custom kernel:
- Input: GEMM output [1024, 8192]
- We need linear weights' bias too (from nn.Linear)
- GroupNorm running stats/parameters
- Bias parameter [1, 8192, 1, 1]
- Output: [1024, 8192, 1, 1] or [1024, 8192]
But the GroupNorm parameters are learned during training - we need to use the trained values from the reference model. Since the problem says "numerically equivalent", we need the same parameters.
The challenge is: can we load the parameters in ModelNew? Yes, we can. The constructor takes the same arguments, so we can initialize the same submodules and extract their weights, or directly create parameters.
Actually, looking more carefully at the problem, we need to match the output numerically. The easiest way is to still create internal nn.Linear and nn.GroupNorm in ModelNew to get their initialized parameters, then use them in the custom kernel.
Wait, but if we create the same modules, we'll have the same parameter initialization (which is what we need for correctness if the test loads state dict or uses the same random seed). The test probably initializes Model with random weights, then calls it, and likewise for ModelNew. So we need ModelNew to initialize identically.
Let me re-read the reference: `get_init_inputs()` returns `[in_features, out_features, num_groups, bias_shape]`. And `get_inputs()` returns `[torch.rand(batch_size, in_features)]`.
So the test does:
```python
model = Model(*get_init_inputs())
model_new = ModelNew(*get_init_inputs())
x = get_inputs()
out_ref = model(x)
out_new = model_new(x)
```
For the outputs to match with randomly initialized weights, we need the same initialization. PyTorch's nn.Linear uses Kaiming uniform for weights and uniform for bias. nn.GroupNorm initializes weight=1 and bias=0.
In ModelNew.__init__, if we create the same nn.Linear and nn.GroupNorm, we should get the same initialization. But PyTorch's initialization depends on the random seed, and creating the modules at different times could affect this... Actually, no. The problem says "numerically equivalent outputs", and since `get_init_inputs()` will be called once for each model, if we use the same torch.randn for bias and create Linear the same way, it should match. But wait - if the test runs Model and ModelNew with the same random seed before creating them, they'd both get the same initialization. The simplest approach is to use the same PyTorch modules in ModelNew, but instead of calling them sequentially, fuse the operations.
Actually, I think the safest approach is:
1. In ModelNew.__init__, create the same nn.Linear and nn.GroupNorm
2. In forward, call linear and group_norm as before
3. Then call a fused min+bias kernel, or just do it in PyTorch
But that won't give much speedup. The real speedup comes from fusion. Let me think about what we can fuse.
The key insight: after GroupNorm, we take min across dim=1, then broadcast-add bias. The min is the bottleneck of the memory reduction.
If we do a custom kernel that does:
1. Compute GEMM output (or take it from cuBLAS)
2. For each batch item, compute group means and vars
3. Normalize, apply gamma/beta, track min
4. Add bias and write output
This avoids materializing the full GroupNorm output.
But computing group means and vars requires reading the GEMM output. For GEMM output of [1024, 8192], we can tile this.
Actually, since out_features=8192 and num_groups=512, each group has 16 channels. For each batch item i:
- For group g (0 to 511): channels g*16 to g*16+15
- Compute mean of 16 values
- Compute var of 16 values
- For each channel in group: normalize, scale by gamma, shift by beta
- Track global min across all 8192 normalized values
We can do this in a single pass if we process one batch item per block or per warp.
Let me design the kernel:
For each batch item i (0 to 1023):
- A block handles one batch item
- The block loads all 8192 values from GEMM output into shared memory or processes in tiles
- First pass: compute group means (512 groups, 16 channels each)
- Second pass: compute group variances
- Third pass: normalize all values, track minimum
- Then compute output for all channels: min + bias[j]
Wait, we can optimize this. We know there are 8192 channels = 512 groups × 16 channels. We can compute mean and variance for all groups in parallel within a block.
For H100, we have high occupancy. Let's use a block size of 512 or 1024 threads.
Actually, 8192 is a nice power of 2. Let's use 1024 threads per block to process each batch item. Each thread handles 8 channels.
For group stats:
- Group 0: channels 0-15, group 1: 16-31, etc.
- Thread 0 handles channels 0, 1024, 2048, ..., 7168 (8 values from different groups!)
- This is bad for coalescing and group locality.
Better: each warp handles consecutive channels.
- Block has 1024 threads = 32 warps
- Each warp handles 256 consecutive channels (but groups are only 16 wide)
- Actually, let's have each thread handle channels tid, tid+1024, tid+2048, ..., tid+7168 for the first batch item.
Wait, for coalescing: if thread i reads channel i, i+1024, etc., these are strided by 1024. With 1024 threads, consecutive threads read consecutive channels. So if we do:
- Thread 0: channels 0, 1024, 2048, 3072, 4096, 5120, 6144, 7168
- Thread 1: channels 1, 1025, 2049, ...
- But channels 0-15 are in group 0. Threads 0-15 each read one channel from group 0. Then they each read one channel from group 64 (channels 1024-1039). This interleaves groups.
For group reduction, we need all threads in a group to communicate. So if group 0 has channels 0-15, we need threads 0-15 to reduce their values for group 0. This is workable with warp shuffles and shared memory.
Actually, let me think about a simpler approach. We have 512 groups, 16 channels each. We can use 512 threads, where thread g processes all 16 channels of group g. But then we only use 512 threads, not enough to saturate H100 SMs for 1024 batch items.
We have 1024 batch items. If we use 512 threads per block and 1024 blocks, that's 512K threads total, which is good.
In this approach:
- Block b handles batch item b
- Thread g (0 to 511) handles group g, channels g*16 to g*16+15
- Each thread reads 16 consecutive values from the GEMM output
- Computes mean, variance, normalizes, tracks local min
- Then writes out 16 values: min_val + bias[channel]
Wait! The output is the same min_val for ALL channels in the batch item, just with different bias added. So:
output[i, j] = global_min[i] + bias[j]
And since we only have 8192 channels, and bias has shape [1, 8192, 1, 1], bias[j] is just bias[0, j, 0, 0].
So actually, we don't need to write different normalized values per channel! We just need the global min. The output is literally `global_min[i] + bias[j]`.
This is HUGE for optimization. Since the final result doesn't depend on individual normalized channels (only their min), we can:
1. Compute GEMM
2. Compute group means and vars
3. Normalize and find global min
4. Write output[i, j] = global_min[i] + bias[j]
And step 4 is just a broadcasted add! We can write a very simple kernel for this.
Wait, but does the min depend on the normalization? Yes. GN(x)[i, k] = γ_k * (x[i,k] - μ_g) / sqrt(σ_g^2 + ε) + β_k, where g is the group containing channel k.
The min across k of this depends on all the group stats and per-channel affine params. So we can't avoid computing the normalized values at least to find the min.
But once we have the min, the output is very simple. And importantly, we don't need to materialize ALL normalized values - we can stream through them and keep only the running min.
Let me design the kernel more carefully:
Kernel 1: cuBLAS GEMM to compute Y = X @ W^T + b_linear
Kernel 2: For each batch item, compute group norms and find min
- Block per batch item (1024 blocks)
- Threads per block: let's say 256 or 512
- Each thread processes a subset of channels
- For each group of 16 channels:
- Read 16 values
- Compute group mean and variance (requires reduction within the group)
- Normalize, apply gamma and beta
- Update running min
- After all groups are processed, reduce min across all threads
- Write output: for each channel j, out[i, j] = global_min + bias[j]
Wait, but we need to actually write the output. The output shape is [1024, 8192, 1, 1]. We could write it as [1024, 8192] and reshape.
For the output writing: once we have global_min for batch item i, we need to write 8192 values. This could be done by the same block after computing the min.
But actually, since each block already has all the threads, we can have each thread write N/num_threads values.
Let me think about efficiency:
- We need to read 8192 floats per batch item for the GEMM output
- We need to read 8192 gammas and betas
- We need to write 8192 output floats
- Total memory traffic: ~16KB read + 16KB read (gamma/beta) + 16KB write + some for weights/bias of linear
The GEMM itself will be the dominant cost. After that, our custom kernel for GN-min-bias is memory bandwidth bound.
For 1024 batch items:
- Read GEMM output: 1024 * 8192 * 4 = 32 MB
- Read gamma/beta: 8192 * 4 * 2 = 64 KB (cached)
- Write output: 1024 * 8192 * 4 = 32 MB
- Total: ~64 MB at 3.35 TB/s = 0.019 ms, negligible
So the main cost is the GEMM. Can we speed up the GEMM?
For M=1024, N=8192, K=8192 in fp32:
- This is 137 GFLOPs
- At 67 TFLOPS: 2.04 ms theoretical
- cuBLAS should be very efficient for this size
If PyTorch uses cuBLAS already (which it does), our custom kernel won't beat cuBLAS on the GEMM. But we can win by:
1. Fusing the GroupNorm + min + bias into a single kernel, avoiding intermediate memory traffic
2. Potentially fusing the Linear bias with the GEMM output (cuBLAS already does this)
Wait, PyTorch's nn.Linear followed by GroupNorm might not be fused. The sequence is:
1. cuBLAS GEMM + bias: writes [1024, 8192]
2. PyTorch GroupNorm: reads [1024, 8192], writes [1024, 8192]
3. Min: reads [1024, 8192], writes [1024, 1]
4. Bias add: reads [1024, 1] and bias, writes [1024, 8192, 1, 1]
Total intermediate memory: 32 + 32 + 0.004 + 32 ≈ 96 MB of intermediate traffic
Our approach:
1. cuBLAS GEMM + bias: writes [1024, 8192] (32 MB)
2. Custom kernel: reads [1024, 8192], gamma/beta, bias; writes [1024, 8192] (64 MB + tiny)
3. Total: 96 MB vs reference which is about 128-160 MB
Hmm, the difference might not be huge. But PyTorch's GroupNorm is probably not perfectly optimized, and there might be additional overheads.
Actually wait - can we fuse step 1 and 2? In theory yes, with CUTLASS or a custom GEMM. But for M=1024, K=8192, N=8192, implementing a custom GEMM that beats cuBLAS is very hard.
Actually, let me reconsider. Maybe we can avoid the full GEMM output write! After cuBLAS computes GEMM, it writes to global memory. Then our kernel reads it. We can't easily fuse these without CUTLASS.
But what if we process the GEMM output tile by tile? The problem is we need the full row (all 8192 channels of a batch item) to compute group stats.
Hmm, actually for this specific case, there might be a smarter approach using CUTLASS epilogue fusion. But CUTLASS is complex.
Let me think if we can formulate this as a single batched operation or use tensor cores.
Actually, for the H100, I should check if we can use Warp Group Matrix Multiply (WGMMA) in Hopper. But that's quite advanced.
Let me start with a simpler approach that should still give a decent speedup:
1. Call at::linear (or torch.matmul + bias) for the GEMM
2. Custom kernel for GroupNorm + min + bias add
The key speedup comes from:
- Fusing the GroupNorm parameters (gamma, beta) directly
- Computing the min while normalizing (no need to write/read intermediate GroupNorm output)
- Broadcasting the min + bias in one step
Let me design the CUDA kernel for step 2.
For each batch item i, and each group g:
Group g contains channels [g*16, g*16+15]
Mean_g = sum_{c in group} y[i, c] / 16
Var_g = sum_{c in group} (y[i, c] - Mean_g)^2 / 16
For each c in group:
normalized = ((y[i, c] - Mean_g) / sqrt(Var_g + eps)) * gamma[c] + beta[c]
min_i = min(min_i, normalized)
Then output[i, c] = min_i + bias[c]
Kernel design:
- Grid: [batch_size] = 1024 blocks
- Block: 512 threads
Why 512 threads? We have 512 groups. If each thread handles one group, that's perfect!
- Thread g processes group g (16 channels)
- Reads 16 channels from input
- Computes mean and variance locally
- Normalizes, tracks local min across its 16 channels
- Writes local min to shared memory
- Block sync
- Reduce 512 local mins to 1 global min (within the block)
- Then each thread writes the output for its assigned channels: global_min + bias[c]
Wait, with 512 threads and 8192 channels, each thread would write 16 output channels. With 512 threads per block and 1024 blocks, total threads = 512K. Each thread reads 16 floats, does some math, writes 16 floats. This is very efficient.
And 1024 blocks should keep all SMs busy on H100 (which has ~132 SMs? No, H100 has 132 or 144 SMs for SXM5 actually).
Let me check: H100 SXM5 has 132 SMs (for non-sparse) or 144? Actually, H100 has 132 SMs for PCIe and 132 for SXM5. wait, let me think: H100 80GB SXM5 has 132 SMs, each with 128 FP32 cores.
Total threads if we use 1024 threads per block * 1024 blocks = 1M threads, which is plenty. But 512 threads per block * 1024 blocks = 512K threads, also fine.
Actually, 512 threads per block gives higher occupancy since H100 allows up to 32 blocks per SM, and up to 2048 threads per SM. With 512 threads/block, we can fit 4 blocks per SM, so 1024 blocks / (132 SMs * 4) = ~2 waves, which is fine.
But there's another consideration: coalesced memory access. If thread g reads channels g*16 to g*16+15, these are consecutive in memory! So thread 0 reads channels 0-15, thread 1 reads 16-31, etc. This is perfectly coalesced.
Actually, wait. In the PyTorch linear output, channels are stored consecutively: y[i, 0], y[i, 1], ..., y[i, 8191]. So if thread g reads starting at g*16, consecutive threads read consecutive 16-float blocks. This is excellent coalescing.
For 512 threads per block:
- Thread 0: reads input[i, 0:16]
- Thread 1: reads input[i, 16:32]
- ...
- Thread 511: reads input[i, 8160:8176]
Wait, that's only 512*16 = 8192 channels! Perfect.
Actually, the block size should be 512. But what if I want to use fewer threads and have each thread handle multiple groups? Let's see: if we use 256 threads, each thread handles 2 groups (32 channels). Still good. But 512 is nice because it matches the number of groups.
With 512 threads per block:
- Each thread loads 16 floats = 64 bytes (one cacheline on H100 is 128 bytes, so each thread loads half a cacheline - but consecutive threads load consecutive cachelines, perfect)
- Since 16 floats is small, we could use float4 loads for 4 floats at a time.
Let me write the kernel:
```cuda
template <typename T>
__global__ void fused_gn_min_bias_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
const T* __restrict__ bias,
T* __restrict__ output,
int batch_size,
int num_channels,
int num_groups,
float eps
) {
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
int channels_per_group = num_channels / num_groups;
int group_idx = tid; // each thread handles one group
const T* batch_input = input + batch_idx * num_channels;
T* batch_output = output + batch_idx * num_channels;
float local_min = INFINITY;
if (group_idx < num_groups) {
int start_c = group_idx * channels_per_group;
int end_c = start_c + channels_per_group;
// Compute mean
float sum = 0.0f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
sum += (float)batch_input[c];
}
float mean = sum / channels_per_group;
// Compute variance
float sq_sum = 0.0f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
float diff = (float)batch_input[c] - mean;
sq_sum += diff * diff;
}
float var = sq_sum / channels_per_group;
float inv_std = rsqrtf(var + eps);
// Normalize and find local min
#pragma unroll
for (int c = start_c; c < end_c; c++) {
float normalized = ((float)batch_input[c] - mean) * inv_std;
normalized = normalized * (float)gamma[c] + (float)beta[c];
local_min = fminf(local_min, normalized);
}
}
// Reduce min across all threads in block
__shared__ float shared_min[512];
shared_min[tid] = local_min;
__syncthreads();
// Butterfly reduction
#pragma unroll
for (int stride = 256; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_min[tid] = fminf(shared_min[tid], shared_min[tid + stride]);
}
__syncthreads();
}
float global_min = shared_min[0]; // broadcast from shared memory
// Write output
if (group_idx < num_groups) {
int start_c = group_idx * channels_per_group;
int end_c = start_c + channels_per_group;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
batch_output[c] = (T)(global_min + (float)bias[c]);
}
}
}
```
Wait, there are issues:
1. The output shape in the reference is [1024, 8192, 1, 1], not [1024, 8192]. I need to check if the test checks exact shape or just values.
2. bias has shape (1, out_features, 1, 1), so bias[c] might need different indexing.
3. The reduction uses shared memory and needs careful handling.
For the output shape: the reference returns x with shape [1024, 8192, 1, 1] after broadcasting add. In PyTorch, if our output is [1024, 8192] or [1024, 8192, 1, 1], the values are the same, just different shape. But the test probably checks shape equality too! I need to make sure output has the right shape.
Let me look at the reference again:
```python
x = x + self.bias # self.bias is (1, out_features, 1, 1)
# Result: [1024, 8192] + [1, 8192, 1, 1] = [1024, 8192, 1, 1] (via broadcasting)
```
So the output shape is [1024, 8192, 1, 1]. I should make sure our kernel produces this shape. We can either:
- Write the custom kernel to output [1024, 8192] and then `.unsqueeze(-1).unsqueeze(-1)`
- Or write directly to [1024, 8192, 1, 1]
For simplicity, I'll output [1024, 8192] and then `.view(batch_size, out_features, 1, 1)` or use `unsqueeze`.
Actually, looking at the bias shape: `bias_shape = (1, out_features, 1, 1)`. So `self.bias` is [1, 8192, 1, 1]. When added to [1024, 1], the result is [1024, 8192, 1, 1].
For our custom kernel, we can read bias as a flat array of 8192 elements and write output with strides for the [1024, 8192, 1, 1] shape. The output will be contiguous in the first two dims and singletons in the last two, so effectively it's just [batch, channel] stored contiguously with total size batch * channel. The shape is just metadata.
Actually, in PyTorch, tensor [1024, 8192, 1, 1] with shape ().numel() = 1024*8192, and it's contiguous. So if we write a flat buffer of 1024*8192 elements, we can just `.view(1024, 8192, 1, 1)` and it's the same memory and values.
So my kernel can write a flat output and we reshape it.
Now, about the indexing of bias. The bias is passed as `bias.data_ptr<T>()`. Since it's [1, 8192, 1, 1], channel c is at `bias[c]` if we assume contiguous memory with the channel stride being 1. Let me verify: for shape [1, 8192, 1, 1], strides are [8192, 1, 1, 1]. So yes, bias[c] is at offset c.
But it's safer to use `bias.data_ptr<T>()[c]` with c < 8192. The actual channel index in the bias tensor is at dimension 1, but since dims 0, 2, 3 are all size 1, the linear index is just the channel index.
Wait, no. For tensor indexing: element at [0, c, 0, 0] has offset = 0 * 8192 + c * 1 + 0 * 1 + 0 * 1 = c. So yes, offset c.
For output: element at [b, c, 0, 0] has offset = b * 8192 + c * 1 + 0 * 1 + 0 * 1 = b * 8192 + c. So a flat index of b * 8192 + c.
Great, so the kernel can use simple indexing.
Let me refine the kernel. I need to be more careful about the shared memory reduction:
```cuda
__shared__ float shared_min[512]; // need to match block size
```
Actually, I should define `const int BLOCK_SIZE = 512;` and use that.
But what if I want to make it more general? For now, since num_groups=512 is fixed by the constructor args (but actually the kernel should work for the given sizes), I can hardcode or template.
Actually, looking at the problem, the specific sizes are given but the constructor is general. However, for a kernel launch, we need to know block sizes at compile time or dynamically. Since we're using load_inline, we can just make it work for any block size up to 1024.
Actually, for flexibility, I can write:
```cuda
extern __shared__ float shared_mem[];
```
But it's easier to just use a fixed max size.
Let me use a fixed shared memory array of 1024 floats, which is small (4KB).
```cuda
template <typename T>
__global__ void fused_gn_min_bias_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
const T* __restrict__ bias,
T* __restrict__ output,
int batch_size,
int num_channels,
int num_groups,
float eps
) {
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
int block_dim = blockDim.x;
int channels_per_group = num_channels / num_groups;
const T* batch_input = input + batch_idx * num_channels;
T* batch_output = output + batch_idx * num_channels;
int groups_per_thread = (num_groups + block_dim - 1) / block_dim;
float local_min = 1e30f;
for (int g_offset = 0; g_offset < groups_per_thread; g_offset++) {
int group_idx = tid + g_offset * block_dim;
if (group_idx >= num_groups) continue;
int start_c = group_idx * channels_per_group;
int end_c = start_c + channels_per_group;
float sum = 0.0f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
sum += (float)batch_input[c];
}
float mean = sum / channels_per_group;
float sq_sum = 0.0f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
float diff = (float)batch_input[c] - mean;
sq_sum += diff * diff;
}
float var = sq_sum / channels_per_group;
float inv_std = rsqrtf(var + eps);
#pragma unroll
for (int c = start_c; c < end_c; c++) {
float normalized = ((float)batch_input[c] - mean) * inv_std;
normalized = normalized * (float)gamma[c] + (float)beta[c];
local_min = fminf(local_min, normalized);
}
}
__shared__ float shared_min[1024];
shared_min[tid] = local_min;
__syncthreads();
for (int stride = block_dim / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_min[tid] = fminf(shared_min[tid], shared_min[tid + stride]);
}
__syncthreads();
}
float global_min = shared_min[0];
// Write output - each thread writes some channels
int channels_per_thread = (num_channels + block_dim - 1) / block_dim;
for (int c_offset = 0; c_offset < channels_per_thread; c_offset++) {
int c = tid + c_offset * block_dim;
if (c < num_channels) {
batch_output[c] = (T)(global_min + (float)bias[c]);
}
}
}
```
Wait, there's an issue with block_dim not being a power of 2. But we can use a serial reduction after. Actually, PyTorch GroupNorm default eps is 1e-5.
But I'm duplicating work. Each thread first processes some groups to compute local_min, then writes some channels. The two mappings don't match. For example, if block_dim=512, thread 0 processes group 0 (channels 0-15) and then writes channels 0, 512, 1024, etc.
For the write phase, if we want coalesced access, we should have thread tid write channels tid, tid+block_dim, tid+2*block_dim, etc. But actually, for the output, since every channel gets the same value `global_min + bias[c]`, and bias is read from global memory, consecutive threads should write consecutive channels for coalescing.
Wait, `bias[c]` is different for each c, so we need to read it. If thread 0 writes channels 0, 512, 1024... that's strided and bad.
Better to have each thread write a contiguous block of channels. For block_dim=512 and num_channels=8192, each thread writes 16 consecutive channels. So thread 0 writes 0-15, thread 1 writes 16-31, etc.
But in this case, the thread's `tid` maps to both group processing and output writing, which works perfectly if block_dim = num_groups = 512.
So let's use block_dim = num_groups = 512 exactly. Then each thread tid:
1. Processes group tid (channels tid*16 to tid*16+15)
2. Writes channels tid*16 to tid*16+15 in the output
This is perfectly coalesced for both read and write!
```cuda
template <typename T>
__global__ void fused_gn_min_bias_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
const T* __restrict__ bias,
T* __restrict__ output,
int num_channels,
int num_groups,
float eps
) {
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
int channels_per_group = num_channels / num_groups;
int start_c = tid * channels_per_group;
int end_c = start_c + channels_per_group;
const T* batch_input = input + batch_idx * num_channels;
T* batch_output = output + batch_idx * num_channels;
// Compute mean
float sum = 0.0f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
sum += (float)batch_input[c];
}
float mean = sum / channels_per_group;
// Compute variance
float sq_sum = 0.0f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
float diff = (float)batch_input[c] - mean;
sq_sum += diff * diff;
}
float var = sq_sum / channels_per_group;
float inv_std = rsqrtf(var + eps);
// Normalize and find local min
float local_min = 1e30f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
float normalized = ((float)batch_input[c] - mean) * inv_std;
normalized = normalized * (float)gamma[c] + (float)beta[c];
local_min = fminf(local_min, normalized);
}
// Reduce min across all threads in block
__shared__ float shared_min[512];
shared_min[tid] = local_min;
__syncthreads();
#pragma unroll
for (int stride = 256; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_min[tid] = fminf(shared_min[tid], shared_min[tid + stride]);
}
__syncthreads();
}
float global_min = shared_min[0];
// Write output
#pragma unroll
for (int c = start_c; c < end_c; c++) {
batch_output[c] = (T)(global_min + (float)bias[c]);
}
}
```
This requires block size exactly 512. If num_groups changes, we'd need a different kernel. But for our problem, num_groups = 512.
Actually, should I make it more general? The problem says the constructor accepts the same arguments, but the inputs in get_init_inputs are fixed. And the test might instantiate with the same fixed values. Let me make the kernel general by using dynamic shared memory or a template.
To be safe, I can write a slightly more general version, but for performance, hardcoding 512 is fine.
Wait, I just realized: `num_groups` comes from the constructor, so it could theoretically be different if the test uses different init inputs. But `get_init_inputs` returns fixed values. I think it's safe to hardcode for the given sizes, but let me keep it somewhat general.
Actually, since we're using `load_inline`, we can just pass `num_groups` as a parameter and use it in the kernel. But shared memory size needs to be known. We can use `extern __shared__ float shared_min[];` and pass the shared memory size at launch.
Or, simpler: we can use a fixed maximum shared memory. Since num_groups is at most out_features (8192), and shared memory per block max is 48KB or so, we can allocate 8192 * 4 = 32KB max. That's fine.
Let me write a general version:
```cuda
template <typename T>
__global__ void fused_gn_min_bias_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
const T* __restrict__ bias,
T* __restrict__ output,
int num_channels,
int num_groups,
float eps
) {
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
int block_dim = blockDim.x;
int channels_per_group = num_channels / num_groups;
const T* batch_input = input + batch_idx * num_channels;
T* batch_output = output + batch_idx * num_channels;
int groups_per_thread = (num_groups + block_dim - 1) / block_dim;
float local_min = 1e30f;
for (int g = 0; g < groups_per_thread; g++) {
int group_idx = tid + g * block_dim;
if (group_idx >= num_groups) continue;
int start_c = group_idx * channels_per_group;
int end_c = start_c + channels_per_group;
float sum = 0.0f;
for (int c = start_c; c < end_c; c++) {
sum += (float)batch_input[c];
}
float mean = sum / channels_per_group;
float sq_sum = 0.0f;
for (int c = start_c; c < end_c; c++) {
float diff = (float)batch_input[c] - mean;
sq_sum += diff * diff;
}
float var = sq_sum / channels_per_group;
float inv_std = rsqrtf(var + eps);
for (int c = start_c; c < end_c; c++) {
float normalized = ((float)batch_input[c] - mean) * inv_std;
normalized = normalized * (float)gamma[c] + (float)beta[c];
local_min = fminf(local_min, normalized);
}
}
extern __shared__ float shared_min[];
shared_min[tid] = local_min;
__syncthreads();
for (int stride = block_dim / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_min[tid] = fminf(shared_min[tid], shared_min[tid + stride]);
}
__syncthreads();
}
float global_min = shared_min[0];
int channels_per_thread = (num_channels + block_dim - 1) / block_dim;
for (int c_offset = 0; c_offset < channels_per_thread; c_offset++) {
int c = tid + c_offset * block_dim;
if (c < num_channels) {
batch_output[c] = (T)(global_min + (float)bias[c]);
}
}
}
```
But this general version has strided writes. For our specific case with block_dim = num_groups = 512, we get contiguous writes by each thread.
Actually wait: if block_dim = 512 and num_channels = 8192, channels_per_thread = 16. Thread tid writes channels tid, tid+512, tid+1024, ..., tid+7680. These are strided by 512! That's terrible for coalescing.
For coalesced writes, we want thread tid to write contiguous channels. So for 512 threads and 8192 channels, threads 0-511 should write channels 0-15, 16-31, etc. But with `c = tid + c_offset * block_dim`, thread 0 writes 0, 512, 1024, etc.
To get contiguous channels: `c = tid * channels_per_thread + c_offset` where channels_per_thread = 16.
So thread 0 writes 0, 1, 2, ..., 15; thread 1 writes 16, 17, ..., 31.
For the group processing, we also want thread tid to process group tid. That maps perfectly: thread 0 processes group 0 (channels 0-15) and writes channels 0-15. Thread 1 processes group 1 (channels 16-31) and writes channels 16-31.
This is perfect. But it only works when block_dim == num_groups, and each thread handles exactly one group.
So for our specific case where num_groups = 512:
- launch with 512 threads per block
- each thread handles exactly one group
- each thread writes exactly channels_per_group = 16 channels
This is optimal and simple.
But what about when the numbers don't divide evenly? Out_features = 8192, num_groups = 512, 8192/512 = 16. That's exact.
For the general case, if num_groups > block_dim, we'd need each thread to handle multiple groups. If num_groups < block_dim, some threads would be idle. For our specific case, 512 == 512, perfect.
I can write a kernel that works when block_dim >= num_groups, but to be safe for the specific test, I'll hardcode block size 512 or just use 512 threads and assert num_groups <= 512.
Actually, I can also support block_dim < num_groups by having each thread process multiple groups. For the write phase, if block_dim < num_groups and num_channels / block_dim > channels_per_group, each thread writes more channels than one group. That's fine.
But let's target the specific case. Our test will use num_groups=512. I'll launch with 512 threads.
Let me now think about writing the full Python file.
```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>
#include <cmath>
template <typename T>
__global__ void fused_gn_min_bias_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
const T* __restrict__ bias,
T* __restrict__ output,
int num_channels,
int num_groups,
float eps
) {
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
int channels_per_group = num_channels / num_groups;
int start_c = tid * channels_per_group;
int end_c = start_c + channels_per_group;
const T* batch_input = input + batch_idx * num_channels;
T* batch_output = output + batch_idx * num_channels;
float sum = 0.0f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
sum += static_cast<float>(batch_input[c]);
}
float mean = sum / static_cast<float>(channels_per_group);
float sq_sum = 0.0f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
float diff = static_cast<float>(batch_input[c]) - mean;
sq_sum += diff * diff;
}
float var = sq_sum / static_cast<float>(channels_per_group);
float inv_std = rsqrtf(var + eps);
float local_min = 1e30f;
#pragma unroll
for (int c = start_c; c < end_c; c++) {
float normalized = (static_cast<float>(batch_input[c]) - mean) * inv_std;
normalized = normalized * static_cast<float>(gamma[c]) + static_cast<float>(beta[c]);
local_min = fminf(local_min, normalized);
}
__shared__ float shared_min[512];
shared_min[tid] = local_min;
__syncthreads();
#pragma unroll
for (int stride = 256; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_min[tid] = fminf(shared_min[tid], shared_min[tid + stride]);
}
__syncthreads();
}
float global_min = shared_min[0];
#pragma unroll
for (int c = start_c; c < end_c; c++) {
batch_output[c] = static_cast<T>(global_min + static_cast<float>(bias[c]));
}
}
torch::Tensor fused_gn_min_bias(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, torch::Tensor bias, int num_groups, float eps) {
int batch_size = input.size(0);
int num_channels = input.size(1);
auto output = torch::empty({batch_size, num_channels}, input.options());
const int block_size = num_groups;
const int grid_size = batch_size;
AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "fused_gn_min_bias", [&] {
fused_gn_min_bias_kernel<scalar_t><<<grid_size, block_size>>>(
input.data_ptr<scalar_t>(),
gamma.data_ptr<scalar_t>(),
beta.data_ptr<scalar_t>(),
bias.data_ptr<scalar_t>(),
output.data_ptr<scalar_t>(),
num_channels,
num_groups,
eps
);
});
return output;
}
'''
cpp_source = '''
torch::Tensor fused_gn_min_bias(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, torch::Tensor bias, int num_groups, float eps);
'''
fused_op = load_inline(
name='fused_gn_min_bias',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['fused_gn_min_bias'],
extra_cuda_cflags=['-O3'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, in_features, out_features, num_groups, bias_shape):
super(ModelNew, self).__init__()
self.gemm = nn.Linear(in_features, out_features)
self.group_norm = nn.GroupNorm(num_groups, out_features)
self.bias = nn.Parameter(torch.randn(bias_shape))
self.num_groups = num_groups
def forward(self, x):
x = self.gemm(x)
# x is [batch, out_features]
gamma = self.group_norm.weight
beta = self.group_norm.bias
# Flatten bias to [out_features]
bias_flat = self.bias.view(-1)
x = fused_op.fused_gn_min_bias(x, gamma, beta, bias_flat, self.num_groups, self.group_norm.eps)
# x is [batch, out_features]
# Reshape to match reference: [batch, out_features, 1, 1]
x = x.unsqueeze(-1).unsqueeze(-1)
return x
```
Wait, there are several issues:
1. The bias in GroupNorm is `self.group_norm.bias` which has shape [out_features]. But the reference model also has `self.bias` with shape [1, out_features, 1, 1]. I named the flat bias `bias_flat` and passed it. But in the kernel, I named it `bias`. Wait, the parameter is `self.bias` which is [1, 8192, 1, 1]. The GroupNorm bias is `self.group_norm.bias` which is [8192].
Wait, I need to be careful:
- The reference's `self.bias` is the bias added AFTER the min. Shape: [1, out_features, 1, 1]
- `self.group_norm` has its own weight and bias parameters (gamma and beta), both shape [out_features]
In my code:
- `gamma = self.group_norm.weight` (shape [8192])
- `beta = self.group_norm.bias` (shape [8192]) -- this is the GN affine shift
- `bias_flat = self.bias.view(-1)` (shape [8192]) -- this is the final bias addition
So in the kernel:
- `gamma[c]` is self.group_norm.weight[c]
- `beta[c]` is self.group_norm.bias[c]
- `bias[c]` is self.bias[0, c, 0, 0] -- the final bias
This is correct.
But wait, in PyTorch's nn.GroupNorm, the default affine is True, so it has weight and bias. And eps is 1e-5 by default.
Another issue: input to our kernel is the output of linear, which is [1024, 8192]. Our kernel expects [batch, channels]. Good.
Now, one important thing: my kernel assumes block_size == num_groups. If I launch with block_size = num_groups, and num_groups is 512, then I need to make sure the shared_min array is large enough. I hardcoded [512], which is fine.
But what if the test uses different num_groups? The get_init_inputs returns num_groups=512 specifically. So block_size=512 is fine.
Actually, I realize I should check whether the output shape needs to match exactly. The reference output is [1024, 8192, 1, 1]. Our output after unsqueeze is also [1024, 8192, 1, 1]. Good.
One more concern: the kernel is written with block_size == num_groups in mind, but I launch with `block_size = num_groups`. If num_groups is not 512 but something else, the shared_min[512] might be too small. I should make it dynamically sized or at least large enough.
Actually, let me just make shared memory dynamically allocated:
```cuda
extern __shared__ float shared_min[];
```
And pass `num_groups * sizeof(float)` as the third launch argument: `<<<grid_size, block_size, num_groups * sizeof(float)>>>`
This is cleaner and general.
Also, I should add some error checking or use __launch_bounds__ to help the compiler. But for now, let's keep it simple.
Wait, there's another issue. If block_size != num_groups, the kernel's indexing `start_c = tid * channels_per_group` is wrong. It assumes thread tid processes group tid, which requires tid < num_groups and block_size = num_groups.
So the kernel ONLY works when block_size == num_groups. I should add a comment or assert. Since we control the launch to use block_size = num_groups, this is fine. But maybe I should write it more clearly.
Actually, for safety, I should make the kernel handle block_size != num_groups. But then I lose the coalescing guarantee. For our specific test, block_size == num_groups is guaranteed, so let me just put a note.
Actually, I want to be extra careful. Let me write the kernel to work even if block_size differs, but with a preference for block_size == num_groups. No, the code is cleaner if we just assert it.
Hmm, looking at PyTorch's load_inline, if we dispatch the kernel, compile errors will show up. Let me proceed.
Actually, I just realized there could be an issue with how the min reduction works. In the reference:
```python
x = torch.min(x, dim=1, keepdim=True)[0]
```
Here `x` after GroupNorm has shape [1024, 8192]. The min is taken over dim=1 (channels), so for each batch item, we get one minimum value across all channels.
In our kernel, we compute global_min correctly across all groups/channels. That matches.
Then `x + self.bias` where self.bias is [1, 8192, 1, 1]. Since x is [1024, 1] and bias is [1, 8192, 1, 1], the result is [1024, 8192, 1, 1].
In our kernel, for each batch item i and channel c, we output `global_min[i] + bias[c]`. That's equivalent to broadcasting [1024, 1] to [1024, 8192] and adding [8192]. Correct!
Now let me think about numerical precision. The reference does:
1. Linear: precision of cuBLAS GEMM (TF32 by default on Ampere/Hopper if allowed, but the problem says fp32 precision)
2. GroupNorm: PyTorch's implementation
3. Min: exact
4. Add: exact
If PyTorch uses TF32 for matmul by default, our kernel should too for the GEMM part. We can use `torch.backends.cuda.matmul.allow_tf32` to check, but typically we just call `self.gemm(x)` and PyTorch handles it.
For our custom kernel, we do computations in float32 (casting to float). This should match PyTorch's GroupNorm which also does fp32 computations.
One more thing: in the reference model, `self.group_norm = nn.GroupNorm(num_groups, out_features)`. PyTorch GroupNorm with affine=True (default) has weight initialized to ones and bias initialized to zeros. This is different from the `self.bias` which is initialized with `torch.randn`.
Our ModelNew also creates the same modules, so the initialization should match exactly.
Let me now think about whether we can do even better.
Actually, an important realization: could we avoid the full GEMM output materialization by computing the min during the GEMM?
For M=1024, N=8192, K=8192:
- The output Y[i, j] = sum_k X[i, k] * W[j, k] + b[j]
- For each batch row i, we need all Y[i, 0:8192] to compute group norms and then min
To compute group norms, we need group means and variances. The variance depends on all values in the group. This means we need the full GEMM output (or at least all values) before we can compute any normalized values. We can't stream the GEMM output and immediately compute normalized min without buffering.
Unless... we buffer the GEMM output in shared memory as we compute it. But for 8192 floats per row, that's 32KB, which fits in shared memory per block. The challenge is computing the GEMM itself in a way that writes to shared memory rather than global memory.
This is possible with a custom GEMM kernel where each block computes one row of the output, stores it in shared memory, then does the GroupNorm + min, then writes the final result. This avoids the global memory round-trip for the GEMM output.
For M=1024, we'd have 1024 blocks, each computing one row (1x8192 output). Each block needs to compute:
Y[i, :] = X[i, :] @ W^T + b[:]
Where X[i, :] is [8192] and W^T is [8192, 8192].
Computing one row of output is essentially a matrix-vector dot product for each column. With 8192 columns, and each dot product being [8192] x [8192] = reduction over K.
This is a batched MV or a skinny GEMM. We can have the block's threads cooperate to compute the dot products.
But implementing a custom GEMM that beats cuBLAS for a single case is risky. cuBLAS is highly optimized. The overhead of writing to global memory and reading back might be worth it given the simplicity.
Actually, let me reconsider. For M=1024, N=8192, K=8192:
- cuBLAS time: ~1-2ms
- Our custom kernel time: ~0.02ms
- Total reference: ~1.5-2.5ms (including GN, min, add)
- Our fused approach: ~1.0-2.0ms
The speedup from fusing the epilogue might be 10-20%. Is that enough?
Actually, PyTorch's GroupNorm on [1024, 8192] with 512 groups is probably not super fast because:
- It's a somewhat general implementation
- It might launch multiple kernels
- It might use atomic operations or multiple reductions
Let me think about what PyTorch's GroupNorm does:
1. Compute mean per group (512 groups per batch item × 1024 batch items = 524,288 group means)
2. Compute variance per group
3. Normalize, scale, shift
4. Write output
This is all memory bandwidth bound. And PyTorch likely has an optimized CUDA kernel for it. The fused kernel should still be faster because we avoid one memory round-trip and the min is integrated.
But to truly get a big win, we might want to fuse the GEMM with the rest. Let me think if we can do this with CUTLASS or a custom kernel.
For the custom row-wise GEMM + GroupNorm:
- Block size: let's say 512 threads
- Each block computes one batch row's output
- Shared memory: 32KB for the row buffer (8192 floats)
- Each thread computes some columns' dot products
For column j, Y[j] = sum_k X[k] * W[j, k]. We can have each thread handle 16 columns (since 8192/512 = 16).
For each thread tid, it handles columns tid*16 to tid*16+15.
To compute these, for each k from 0 to 8191, we need X[k] and W[j, k] for each j in the thread's range.
This is like a GEMM where we tile over K. For 512 threads and K=8192:
- We can load X[k] into shared memory or a register (broadcast to all threads)
- Each thread loads W[j, k] for its 16 columns
But W is stored as [out_features, in_features], so W[j, k] for fixed j and varying k is consecutive. For 16 values of j, W[tid*16, k], W[tid*16+1, k], ... these are strided by in_features=8192. So loading W for 16 different j's at the same k involves 16 strided loads, which is not coalesced.
Alternatively, if we transpose W to [in_features, out_features], then W[k, j] for consecutive j is consecutive in memory. But nn.Linear stores weight as [out_features, in_features]. We can transpose it once in __init__.
In our ModelNew, we can store W_t = self.gemm.weight.t().contiguous(). The bias is separate.
Then for computing row i:
- For each k tile:
- Load X[i, k:k+tile] into shared memory (or registers)
- Each thread loads W_t[k, j:j+16] for its columns
- Compute partial sums
But this is getting complicated. And cuBLAS is very good at this GEMM. The memory traffic savings from avoiding the intermediate write might not justify the potential performance loss from a custom GEMM.
However, let me check if using torch.mm + a custom epilogue is better. Actually, for the H100 with 3.35 TB/s memory bandwidth, writing and reading 32MB takes about 0.01ms (32MB / 3350GB/s = 0.0095ms). So the memory traffic for the GEMM output is negligible compared to the GEMM compute time.
Wait, that calculation seems wrong. 32MB / 3350GB/s = 0.0000095s = 0.0095ms = 9.5 microseconds. Yes, about 10 microseconds. That's indeed negligible compared to 1-2ms for the GEMM.
So the main bottleneck is the GEMM itself, and fusing the epilogue won't help with that. But PyTorch's reference also uses cuBLAS for the GEMM. So our speedup must come from:
1. The GroupNorm + min + add being faster than PyTorch's separate operations
2. Or by having a more efficient overall pipeline
For PyTorch's GroupNorm, it probably uses a well-optimized kernel, but then min is another kernel, and add is another kernel. The sum of these kernels has overhead from multiple launches and some intermediate memory traffic.
The total intermediate memory traffic in the reference:
- Linear output: 32MB write
- GroupNorm output: 32MB write + 32MB read
- Min output: 4KB write + 32MB read
- Final output: 32MB write
Total: ~128MB of intermediate memory + actual compute
Our fused approach:
- Linear output: 32MB write
- Fused kernel: 32MB read + 32MB write
- Total: 64MB of intermediate memory
We save about 64MB of memory traffic. At 3.35 TB/s, that's about 20 microseconds. And we save 2-3 kernel launch overheads (~5-10 microseconds each).
So our theoretical speedup is maybe 30-50 microseconds on a reference that takes ~2ms. That's only 1-2% speedup. That's not very impressive.
Wait, but is PyTorch's GroupNorm for [1024, 8192] with 512 groups really well-optimized and only taking ~0.02ms? Or does it take more?
For PyTorch's GN on this size, the computation per element is:
- Compute mean and var: a few FLOPs
- Normalize: a few FLOPs
Total compute is tiny, it's all memory bandwidth. But kernel launch overhead and the implementation details matter.
Maybe PyTorch's GN implementation is more expensive than we think because:
1. It might compute mean and variance in separate passes (2 reads of input)
2. It might use less than optimal block configurations
3. It might have extra operations for numerical stability
If PyTorch GN takes, say, 0.1ms instead of 0.02ms, then our fusion could save more.
Also, I notice the GroupNorm is operating on a 2D tensor [N, C]. PyTorch's GN kernel might be general and handle ND tensors, which could add some inefficiency for this 2D case.
Let me think about another optimization. Since the final output is just `min_i + bias_j`, we don't actually need to write out all 8192 channels separately if we can somehow avoid it. But the test expects [1024, 8192, 1, 1], so we do need to write 8192 values per row.
One more thought: PyTorch's min(dim=1) on [1024, 8192] launches a reduction kernel. On an H100, a reduction of 8192 elements per row is very fast, but it still requires a kernel launch.
Let me try a different approach. What if we use CUTLASS to fuse the GEMM with the epilogue?
Actually, for the H100, I notice the specs say FP32 TFLOPS is 67, but TF32 Tensor Core is 989 TFLOPS (with sparsity, maybe 494 without?). H100 Tensor Cores support TF32 for FP32 accumulation. If PyTorch enables TF32 by default (which it does on Ampere/Hopper for matmul), the GEMM might be very fast - around 0.3ms at ~500 TFLOPS.
137 GFLOPs / 500 TFLOPS = 0.000274s = 0.274ms. So the GEMM might be ~0.3ms if TF32 is allowed.
Then the remaining operations (GN + min + bias) might be 0.1-0.2ms. If we can fuse them and save memory traffic, maybe we get 0.1ms. That's a 50% speedup of the non-GEMM portion, or about 20% overall.
Actually, I just realized: the problem says "fp32 precision" for numerical equivalence. Does that mean we have to ensure no TF32 is used? Let me check. The tolerance for fp32 is atol=rtol=1e-4. TF32 has about 1e-4 precision (10 mantissa bits vs 24 for fp32), so using TF32 might still pass the 1e-4 tolerance. But to be safe and get exact matching, we should probably match whatever PyTorch does.
If the reference PyTorch model uses `torch.backends.cuda.matmul.allow_tf32 = True` (default), then using TF32 in our GEMM would produce the same result. If the test sets it to False, we need to respect that.
We can't easily control this in our kernel if we call `self.gemm(x)`, because PyTorch's linear layer will follow the global setting. So our result will automatically match.
OK, let me now also consider if we can use a custom kernel for the entire forward pass, including the GEMM, and win that way.
For the custom GEMM approach:
- We have M=1024 rows, N=8192 cols, K=8192 inner dim
- This is C^T = A @ B^T where A is M×K, B is N×K, C is M×N
In CUDA terms, if we launch one block per row (1024 blocks):
- Each block computes C[i, 0:N] = A[i, 0:K] @ B[0:N, 0:K]^T
- Wait, B is weight [N, K], B^T is [K, N]
- C[i, j] = sum_k A[i, k] * B[j, k]
So each block needs to read A[i, :] (8192 floats) and B[:, :] (8192 × 8192 floats, but only columns corresponding to the output).
Actually, B is shared across all blocks. If each block handles one row, all blocks read the entire B matrix over the course of computation. That's 1024 × 8192 × 8192 reads of B = 64GB of B reads! That's terrible for cache.
Actually, each block only reads the columns of B that it needs. If block i computes all N=8192 columns, it reads the entire B matrix (8192 × 8192 floats = 256MB). But there are 1024 blocks, so total B reads = 256MB × 1024 = 256GB. That can't be right...
Wait, let me reconsider. If block 0 computes row 0 of C, it needs:
C[0, j] = sum_k A[0, k] * B[j, k] for j = 0 to 8191
For each j, it reads B[j, 0:8192]. So it reads all of B, but only once per block.
Total B reads across all blocks: 1024 blocks × 8192×8192 floats = 1024 × 256MB = 256GB. That's way too much.
But in practice, this is fine because B stays in L2 cache. Each column j of B is read by all 1024 blocks. If we process one row at a time, B[j, :] is read sequentially by block 0, then block 1, etc. But blocks run concurrently, so they might all be reading different parts of B.
Actually, the issue is that classical GEMM tiling is important. cuBLAS tiles the computation so that each block computes a tile of C (e.g., 128×128), and loads tiles of A and B into shared memory. This maximizes data reuse.
For M=1024, if we tile with 128×128 blocks:
- Grid: 8 × 64 = 512 blocks
- Each block computes a 128×128 tile of C
- Good data reuse of A and B tiles
This is what cuBLAS does, and it's much more efficient.
So I should not attempt a custom GEMM for the main matrix multiply.
But wait - what about using a vector matrix multiply approach? We can compute each C[i, :] using a batched approach or by keeping B in shared memory, but B is too large (256MB) for shared memory (shared mem is 227KB per block).
So definitely use cuBLAS for the GEMM.
Now, let me think about another optimization. What if we use `torch.addmm` or `torch.nn.functional.linear` and then immediately call our cuda kernel? PyTorch might not release the tensor memory, and the GPU can pipeline the operations.
Actually, can we use `torch.utils.dlpack` or something to avoid copies? No, the tensors are already on GPU.
Let me think about whether there's a smarter mathematical formulation.
After linear: Y = XW^T + b_linear
After GN: For each group g, Z[i, c] = γ_c * (Y[i, c] - μ_{i,g}) / σ_{i,g} + β_c
After min: m[i] = min_c Z[i, c]
After bias: O[i, c] = m[i] + bias_c
So O[i, c] = min_c [γ_c * (Y[i, c] - μ_{i,g}) / σ_{i,g} + β_c] + bias_c
Since bias_c is outside the min, and O[i, c] = m[i] + bias_c where m[i] doesn't depend on c.
To compute m[i], for each group g, we compute:
μ_{i,g} = mean(Y[i, g*16:(g+1)*16])
σ_{i,g} = std(Y[i, g*16:(g+1)*16])
Then for each c in group g:
z_c = γ_c * (Y[i, c] - μ_g) / σ_g + β_c
And m[i] = min over all groups g and channels c in g of z_c.
Notice that within a group, the normalization subtracts the same mean and divides by the same std. So z_c = (γ_c / σ_g) * Y[i, c] + (β_c - γ_c * μ_g / σ_g).
Within a group, z_c is an affine function of Y[i, c]. The minimum of an affine function over a set of points occurs at... well, it's just the minimum of the values, not necessarily at any particular point. We still need to compute all z_c and find the min.
OK so there's no mathematical shortcut. We need to compute all normalized values (or at least find the min).
Let me now think about the kernel more carefully and potential issues.
One issue: the shared_min array needs to be sized to the block size. I'll use dynamic shared memory.
Another issue: for very large block sizes, the reduction might need more steps. I'll use a loop.
Wait, I also need to consider that if we allow_tf32 for the GEMM, the reference uses it too, so our GEMM will match. Good.
What about when PyTorch compiles the reference? In eager mode, each op is launched separately. In our case, we have fewer kernels.
Let me also consider if we can pipeline the GEMM and our custom kernel using CUDA streams. But PyTorch's CUDA ops are async already and the CUDA runtime scheduler should handle this. We don't need explicit streams.
Now let me think about the kernel launch configuration more carefully. For our specific case:
- batch_size = 1024
- num_groups = 512
- channels_per_group = 16
- block_size = 512
Each block processes one batch item. Total blocks = 1024.
With 132 SMs and max 32 blocks per SM, we can have up to 4224 blocks in flight. We launch 1024 blocks, so all can be resident.
Occupancy: each block has 512 threads. H100 has 2048 threads per SM max. So 4 blocks per SM, or 2048 threads per SM. That's 100% occupancy if registers and shared memory permit.
Shared memory: we need 512 floats = 2KB per block. H100 has 227KB per SM, so we can fit many blocks.
Registers: the kernel is fairly simple. Each thread probably uses ~32-64 registers. With 512 threads and 64 registers each, that's 32K registers per block. Max is 64K per block, 64K per SM. So 2 blocks per SM at 64 registers/thread. Let's say it fits.
Actually, with float loads and unrolled loops, the register usage might be higher due to keeping values in registers. But 16 iterations of a loop with 512 threads should be fine.
Wait, I have `#pragma unroll` for the loops over channels_per_group=16. The compiler will unroll these. For reading input, computing mean, computing variance, normalizing, and writing output, that's 4×16 = 64 float operations. With unrolling, the compiler might use a lot of registers. But 16 is small enough.
Actually, I don't need `#pragma unroll` for everything. I can let the compiler decide, or unroll just some loops. But 16 iterations is fine.
Let me think about vectorized loads. We can load float4 at a time:
```cuda
const float4* batch_input_vec = reinterpret_cast<const float4*>(batch_input);
```
Since each thread accesses channels start_c to end_c (16 consecutive floats), and 16 is divisible by 4, we can do 4 float4 loads. This increases memory throughput.
Similarly for writing output.
And for gamma, beta, bias: they are accessed at the same channel indices, so float4 loads work there too.
Let me refine the kernel with float4 vectorization:
```cuda
template <typename T>
__global__ void fused_gn_min_bias_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
const T* __restrict__ bias,
T* __restrict__ output,
int num_channels,
int num_groups,
float eps
) {
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
int channels_per_group = num_channels / num_groups;
int start_c = tid * channels_per_group;
const T* batch_input = input + batch_idx * num_channels + start_c;
T* batch_output = output + batch_idx * num_channels + start_c;
const T* gamma_ptr = gamma + start_c;
const T* beta_ptr = beta + start_c;
const T* bias_ptr = bias + start_c;
float sum = 0.0f;
#pragma unroll
for (int c = 0; c < channels_per_group; c++) {
sum += static_cast<float>(batch_input[c]);
}
float mean = sum / static_cast<float>(channels_per_group);
float sq_sum = 0.0f;
#pragma unroll
for (int c = 0; c < channels_per_group; c++) {
float diff = static_cast<float>(batch_input[c]) - mean;
sq_sum += diff * diff;
}
float var = sq_sum / static_cast<float>(channels_per_group);
float inv_std = rsqrtf(var + eps);
float local_min = 1e30f;
#pragma unroll
for (int c = 0; c < channels_per_group; c++) {
float normalized = (static_cast<float>(batch_input[c]) - mean) * inv_std;
normalized = normalized * static_cast<float>(gamma_ptr[c]) + static_cast<float>(beta_ptr[c]);
local_min = fminf(local_min, normalized);
}
extern __shared__ float shared_min[];
shared_min[tid] = local_min;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_min[tid] = fminf(shared_min[tid], shared_min[tid + stride]);
}
__syncthreads();
}
float global_min = shared_min[0];
#pragma unroll
for (int c = 0; c < channels_per_group; c++) {
batch_output[c] = static_cast<T>(global_min + static_cast<float>(bias_ptr[c]));
}
}
```
This should work. The shared memory size is `blockDim.x * sizeof(float)` = 512 * 4 = 2048 bytes.
One issue: the reduction assumes blockDim.x is a power of 2. 512 is, so we're fine. If blockDim.x were not a power of 2, we'd need an extra step. But we launch with blockDim.x = num_groups = 512.
Actually, to be robust, I can handle non-power-of-2 by first doing a check, but it's not necessary for this problem.
Let me now write the complete model and test it.
Wait, one more thing: I need to make sure the input to the kernel is contiguous and the right dtype. If input is fp32, everything is fine. The output should also be fp32.
For bias, gamma, beta: they come from PyTorch parameters and are fp32.
Let me also check: what if num_groups doesn't divide num_channels evenly? In our case, 512 divides 8192. The constructor might not enforce this, but PyTorch GroupNorm requires it.
Actually, looking at the code structure:
```python
fused_op = load_inline(...)
```
This should be at module level. And `ModelNew` uses it.
Let me also think about whether we need to handle the case where the test calls the model multiple times. `load_inline` caches by `name`, so subsequent loads are fast.
Now let me write the code. I'll make sure to handle the output shape correctly.
Actually, I want to check something. In the reference:
```python
x = torch.min(x, dim=1, keepdim=True)[0]
```
If x before min is [1024, 8192], after min it's [1024, 1]. Then:
```python
x = x + self.bias
```
self.bias is [1, 8192, 1, 1]. The broadcasting works because [1024, 1] unsqueezes to match [1, 8192, 1, 1], resulting in [1024, 8192, 1, 1].
In our kernel, we output [batch_size, num_channels] which is [1024, 8192]. Then we unsqueeze. Good.
Now let me think about whether I should also consider fusing in a different way. Since the output is O[i, c] = m[i] + bias_c, and m[i] only depends on the row i, could we compute m[i] using a reduction of some transformed values?
Actually, there's a subtlety. Is there a way to compute the minimum of GN(x) without computing all normalized values explicitly?
For a group g with channels c in [g*16, (g+1)*16], the normalized values are:
z_c = γ_c * (Y_c - μ_g) / σ_g + β_c
= (γ_c / σ_g) * Y_c + (β_c - γ_c * μ_g / σ_g)
Within a group, this is an affine function. But since each channel has different γ_c and β_c, the argmin (which channel has the minimum value) depends on all the values. Since it's an affine function, and the coefficients γ_c / σ_g can be positive or negative... wait, γ_c is initialized to 1.0 in PyTorch's GroupNorm, and during training it can be any positive or negative value.
If all γ_c are positive, then the minimum normalized value in a group corresponds to the minimum Y_c if we ignore the constant term. But with β_c varying, it's not that simple.
We'd still need to check all channels. So no shortcut.
Let me think about another aspect: can we use Tensor Cores or other fast paths for parts of the computation? Probably not for GroupNorm since it's elementwise and reduction, not matrix multiply.
One more optimization: since each thread processes exactly 16 channels, the `#pragma unroll` for loops of size 16 should work well. But is 16 the right number? It's determined by num_channels / num_groups = 8192 / 512 = 16.
Actually, what if we vectorize by processing float4 at a time? With 16 channels, that's 4 float4 operations.
For float4, in CUDA C++:
```cuda
float4 val = reinterpret_cast<const float4*>(batch_input)[0]; // loads 4 floats
```
But we need to be careful about alignment. Since start_c is a multiple of 16 (tid * 16), and 16 * sizeof(float) = 64, which is divisible by 16 (alignment of float4), the pointers should be aligned for float4 loads.
For contiguous memory where thread tid reads at offset tid*16, thread 0 reads at 0, thread 1 at 64, etc. Since 64 is divisible by 16, float4 alignment is satisfied. Good!
Let me write a vectorized version:
```cuda
template <typename T>
__global__ void fused_gn_min_bias_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
const T* __restrict__ bias,
T* __restrict__ output,
int num_channels,
int num_groups,
float eps
) {
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
int channels_per_group = num_channels / num_groups;
int start_c = tid * channels_per_group;
const T* batch_input = input + batch_idx * num_channels + start_c;
T* batch_output = output + batch_idx * num_channels + start_c;
const T* gamma_ptr = gamma + start_c;
const T* beta_ptr = beta + start_c;
const T* bias_ptr = bias + start_c;
float sum = 0.0f;
float sq_sum = 0.0f;
float local_min = 1e30f;
// Process in float4 chunks of 4
#pragma unroll
for (int c = 0; c < channels_per_group; c += 4) {
// Load input
float4 in4 = *reinterpret_cast<const float4*>(batch_input + c);
float4 g4 = *reinterpret_cast<const float4*>(gamma_ptr + c);
float4 b4 = *reinterpret_cast<const float4*>(beta_ptr + c);
sum += in4.x + in4.y + in4.z + in4.w;
// Wait, we can't compute sq_sum until we have mean...
}
// Hmm, we need mean first, then variance, then normalized values.
```
This doesn't vectorize well across the passes because we need mean before variance, and variance before normalization. We could:
1. Load all values into registers first (16 floats = 16 registers)
2. Compute sum
3. Compute variance
4. Compute normalized values and min
With 16 floats in registers, we'd use some registers. But unrolling a loop of 16 is fine without explicit float4.
Actually, let me think about register pressure. With 16 values loaded sequentially:
```cuda
float vals[16];
for (int i = 0; i < 16; i++) vals[i] = batch_input[i];
```
The compiler might put these in registers or shared memory. For 16 floats, it should use registers. But `#pragma unroll` with explicit array declaration might cause issues. Let me just keep the sequential loads.
Actually, for simplicity, let me not over-optimize the inner loops. The memory bandwidth is not the bottleneck here, the kernel launch is. The kernel will be fast enough.
Let me now write the complete code and compile it.
Wait, I also want to check if we should use `aten::cuda::CUDAStreamGuard` or something. With `load_inline`, the kernels use the current CUDA stream, which is what PyTorch uses. So no need for explicit stream handling.
One more thing: our kernel uses `rsqrtf` which is fast on H100. Good.
Let me also make sure we're using float32 for the computation even if input is half/bfloat16. The problem says the model runs at fp32. Actually, the input is `torch.rand(batch_size, in_features)` which defaults to float32. And the modules are float32 too. So we don't need half support, but it's good to have.
Actually, to be safe, I can add a check that the input is contiguous and use `.contiguous()` if needed. But the output of nn.Linear is contiguous, so we should be fine.
OK here's my plan:
1. Write the kernel with block_size = num_groups = 512
2. Call cuBLAS via `self.gemm(x)`
3. Call custom kernel
4. Reshape output
Now let me write the actual code.
One thing I want to verify: the kernel should work when `blockDim.x != num_groups`, but since we launch with `blockDim.x = num_groups`, and `num_groups = 512`, it will work.
Wait, actually I realized that I should pass the correct number of shared memory bytes. For `extern __shared__ float shared_min[];`, I need to pass the size in the kernel launch: `<<<grid, block, shared_mem_size>>>`. The shared_mem_size in bytes is `num_groups * sizeof(float)` for our case.
Alternatively, I can just use a fixed-size static shared memory array `__shared__ float shared_min[512];` and pass `block_size = num_groups`. If num_groups could be larger, this wouldn't work, but it's 512 for our test. To be slightly more general, I can use a template with a max size.
I'll just use `extern __shared__ float shared_min[];` and pass the dynamic shared memory size.
Now, about the `load_inline` call:
```python
fused_op = load_inline(
name='fused_gn_min_bias',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['fused_gn_min_bias'],
extra_cuda_cflags=['-O3', '--use_fast_math'],
verbose=False
)
```
I should check if `-O3` and `--use_fast_math` are appropriate. `--use_fast_math` might make `rsqrtf` less precise but faster. However, for correctness within 1e-4, it should be fine. But to match PyTorch exactly, I might not want `--use_fast_math` since PyTorch's rsqrt might be more precise.
Actually, PyTorch's GroupNorm uses `rsqrt` from ATen, which maps to CUDA's `rsqrtf` or `rsqrt`. For fp32, it should be the same. I'll avoid `--use_fast_math` to be safe with precision.
Let me write the code now:
```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>
template <typename T>
__global__ void fused_gn_min_bias_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
const T* __restrict__ bias,
T* __restrict__ output,
int num_channels,
int num_groups,
float eps
) {
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
int channels_per_group = num_channels / num_groups;
int start_c = tid * channels_per_group;
const T* batch_input = input + batch_idx * num_channels + start_c;
T* batch_output = output + batch_idx * num_channels + start_c;
const T* gamma_ptr = gamma + start_c;
const T* beta_ptr = beta + start_c;
const T* bias_ptr = bias + start_c;
float sum = 0.0f;
#pragma unroll
for (int c = 0; c < channels_per_group; c++) {
sum += static_cast<float>(batch_input[c]);
}
float mean = sum / static_cast<float>(channels_per_group);
float sq_sum = 0.0f;
#pragma unroll
for (int c = 0; c < channels_per_group; c++) {
float diff = static_cast<float>(batch_input[c]) - mean;
sq_sum += diff * diff;
}
float var = sq_sum / static_cast<float>(channels_per_group);
float inv_std = rsqrtf(var + eps);
float local_min = 1e30f;
#pragma unroll
for (int c = 0; c < channels_per_group; c++) {
float normalized = (static_cast<float>(batch_input[c]) - mean) * inv_std;
normalized = normalized * static_cast<float>(gamma_ptr[c]) + static_cast<float>(beta_ptr[c]);
local_min = fminf(local_min, normalized);
}
extern __shared__ float shared_min[];
shared_min[tid] = local_min;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_min[tid] = fminf(shared_min[tid], shared_min[tid + stride]);
}
__syncthreads();
}
float global_min = shared_min[0];
#pragma unroll
for (int c = 0; c < channels_per_group; c++) {
batch_output[c] = static_cast<T>(global_min + static_cast<float>(bias_ptr[c]));
}
}
torch::Tensor fused_gn_min_bias(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, torch::Tensor bias, int num_groups, float eps) {
int batch_size = input.size(0);
int num_channels = input.size(1);
auto output = torch::empty({batch_size, num_channels}, input.options());
const int block_size = num_groups;
const int grid_size = batch_size;
const int shared_mem_size = num_groups * sizeof(float);
AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "fused_gn_min_bias", [&] {
fused_gn_min_bias_kernel<scalar_t><<<grid_size, block_size, shared_mem_size>>>(
input.data_ptr<scalar_t>(),
gamma.data_ptr<scalar_t>(),
beta.data_ptr<scalar_t>(),
bias.data_ptr<scalar_t>(),
output.data_ptr<scalar_t>(),
num_channels,
num_groups,
eps
);
});
return output;
}
"""
CPP_SOURCE = "torch::Tensor fused_gn_min_bias(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, torch::Tensor bias, int num_groups, float eps);"
fused_gn_min_bias_op = load_inline(
name="fused_gn_min_bias",
cpp_sources=CPP_SOURCE,
cuda_sources=CUDA_SOURCE,
functions=["fused_gn_min_bias"],
extra_cuda_cflags=["-O3"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, in_features, out_features, num_groups, bias_shape):
super(ModelNew, self).__init__()
self.gemm = nn.Linear(in_features, out_features)
self.group_norm = nn.GroupNorm(num_groups, out_features)
self.bias = nn.Parameter(torch.randn(bias_shape))
self.num_groups = num_groups
def forward(self, x):
x = self.gemm(x)
gamma = self.group_norm.weight
beta = self.group_norm.bias
bias_flat = self.bias.view(-1)
x = fused_gn_min_bias_op.fused_gn_min_bias(x, gamma, beta, bias_flat, self.num_groups, self.group_norm.eps)
x = x.unsqueeze(-1).unsqueeze(-1)
return x
```
Wait, I need to think about one more thing: the shape of `self.bias`. In the constructor, `bias_shape = (1, out_features, 1, 1)`. So `self.bias` is [1, 8192, 1, 1]. When we call `self.bias.view(-1)`, it becomes [8192]. Good.
But what if bias_shape is different in a general case? The test uses the specific values from `get_init_inputs`. So it should be fine.
Actually, I should check if the element indexing in the kernel for bias is correct. For bias with shape [1, 8192, 1, 1], `bias.view(-1)` gives a contiguous tensor with 8192 elements, and element at channel c is at index c. This matches `bias_ptr[c]` in the kernel.
Similarly for gamma and beta, they are [8192] and contiguous.
One more subtlety: the reference uses `self.group_norm(x)` where x is [1024, 8192]. PyTorch's GroupNorm for 2D input treats it as N, C (with no spatial dimensions, so H=W=1 implicitly). Our kernel also treats input as [N, C] and computes per-group stats across C. This matches.
But wait, in the reference:
```python
x = self.gemm(x) # [1024, 8192]
x = self.group_norm(x) # [1024, 8192]
x = torch.min(x, dim=1, keepdim=True)[0] # [1024, 1]
x = x + self.bias # [1024, 1] + [1, 8192, 1, 1] = [1024, 8192, 1, 1]
```
In PyTorch, when GroupNorm gets a 2D tensor [N, C], it reshapes to [N, C, 1, 1] internally, computes mean/var over [C/group, 1, 1], then normalizes. The output shape is still [N, C] though (or actually, it preserves the input shape). Let me verify.
Actually, looking at PyTorch source, `nn.GroupNorm` calls `F.group_norm` which operates on the input as-is. For input of shape [N, C], the mean/var is computed over the last dimensions... actually, let me check the documentation.
`nn.GroupNorm(num_groups, num_channels)` expects input of shape (N, C, *) where * is any number of additional dimensions. For input [N, C], it treats it as having no spatial dimensions, equivalent to [N, C, 1]. Wait, actually I think it requires at least 2D input and normalizes over the channel groups and all spatial dims. For a 2D tensor [N, C], there are no spatial dims beyond C, so the group normalization divides C into groups and normalizes each group.
The actual computation in PyTorch's group_norm kernel:
- For each sample n and each group g:
- channels = num_channels / num_groups
- The mean is computed over `channels * H * W` elements
- For 2D input [N, C], H=W=1, so it's over `channels` elements
So yes, our kernel matches exactly.
OK, let me compile and test this.