reasoning
Let me analyze the reference model:
1. Linear layer: (batch_size, input_size) @ (input_size, hidden_size)^T + bias
- For the given sizes: (128, 32768) @ (32768, 32768) + bias
- This involves reading input (128 * 32768), weight (32768 * 32768), and bias (32768)
- The output is (128, 32768)
- FLOPs: 128 * 32768 * 32768 * 2 = ~274.9 billion FLOPs
- Memory: ~4.3GB for weight, ~16MB for input, ~1.5MB for output
2. Sigmoid: applied element-wise to (128, 32768)
3. Sum along dim 1: (128, 32768) -> (128, 1)
Given the H100 specs with TF32 Tensor Core at 989 TFLOPS and FP32 at 67 TFLOPS, this is a compute-bound problem (275B FLOPs).
The key optimization here is to fuse all operations into a single custom CUDA kernel using CUTLASS-like approaches or a custom kernel that:
1. Loads the weight matrix
2. Computes each output element = sigmoid(sum(x[j] * W[i,j]) + bias[i])
3. Then sum across each row
Wait, but the sum is over dim 1, which means we sum over the hidden_size dimension. So for each batch element b, output[b] = sum_i sigmoid(sum_j x[b,j] * W[i,j] + bias[i]).
We can fuse all three operations: matmul + sigmoid + sum reduction.
Actually, let me think more carefully. In the reference:
1. x @ W^T gives (batch_size, hidden_size)
2. Add bias to each row
3. sigmoid gives (batch_size, hidden_size)
4. sum over dim 1 gives (batch_size, 1)
The standard cuBLAS approach for matmul -> sigmoid -> sum could be fused. Actually, a custom kernel where each thread/block processes one batch element and computes the final scalar sum by accumulating partial sums.
For batch_size = 128 and hidden_size = 32768:
- Each batch element needs to compute 32768 dot products of length 32768
- Then apply sigmoid to each of the 32768 results
- Then sum all 32768 values
This is huge compute (each batch element: 32768 * 32768 = ~1B ops). 128 batch elements = ~128B ops just for the matmul. Then another ~128M ops for sigmoid+sum.
Wait, for matmul of (128, 32768) with (32768, 32768):
M=128, K=32768, N=32768
The standard approach uses cuBLAS which is highly optimized.
For the reference, PyTorch will call cuBLAS for the linear layer (which includes bias), then elementwise sigmoid, then sum. We can fuse the sigmoid and sum, but the matmul is the bottleneck.
However, we can think about fusing sigmoid and sum into a custom reduction kernel. Or we can use a custom kernel that computes the whole thing at once. Actually, for the matm, using cuBLAS is optimal since it's compute bound. But we can fuse the sigmoid + sum into a single kernel that operates on the output of matmul.
Wait - we can do even better! We can fuse the entire thing if we write a custom matmul kernel that computes the final sum per row directly. But M=128 is very small relative to K=N=32768, so a standard cuBLAS GEMM with the epilogue fused might be best.
Actually, for M=128, K=32768, N=32768, this is a batched vector-matrix multiplication essentially (each row of x is multiplied by W). cuBLAS GEMM should handle this well.
But we can't easily call cuBLAS from a custom kernel in PyTorch load_inline easily. However we can:
1. Still call PyTorch's matmul for the linear part (or use torch.nn.functional.linear which calls cuBLAS)
2. Then call our custom fused kernel for sigmoid + sum reduction
Actually, that won't be faster because matmul itself dominates. Let me calculate:
- M=128, N=32768, K=32768
- Time for cuBLAS GEMM at peak: 2 * M * N * K / FLOPS = 2 * 128 * 32768 * 32768 / (989e12 with TF32, but linear doesn't use tensor cores by default) or / (67e12 FP32)
- FP32: 274.9e9 / 67e12 = ~4.1ms
- With Tensor CoresTF32: 274.9e9 / 989e12 = ~0.28ms (if using tensor cores, but linear in fp32 typically doesn't)
Actually for FP32, cuBLAS should use tensor cores in some modes. The theoretical FP32 TFLOPS is 67, but tensor cores can do TF32 at 989 TFLOPS. However PyTorch linear in fp32 might not use TF32 automatically.
Wait, the output must be fp32 precision (numerically equivalent within 1e-4). tf32 has less precision. So we need fp32 compute, which is 67 TFLOPS.
So around 4ms for the matmul. The sigmoid + sum would be memory-bound at ~32768*128*4 bytes / 3.35 TB/s = ~5 microseconds. So the matmul dominates by far.
If we just fuse sigmoid+sum, we get essentially no speedup.
But wait - can we write a custom kernel that does the linear + sigmoid + sum all in one, and somehow be faster than cuBLAS for this specific shape? For M=128, K=32768, N=32768... this is a very skinny-tall problem where each output row is computed from one input row.
Actually, maybe we can do better by splitting the work differently. Since we're computing one scalar per batch element, and that scalar is sum_i sigmoid(dot(x, W[i]) + bias[i]).
For each output element:
1. For i in [0, N): compute dot(x, W[i]) + bias[i], apply sigmoid, accumulate into sum
This means for each batch element, we iterate over all N rows of W.
With a naive approach, this would be much slower than cuBLAS. But we can process multiple rows of W in parallel with shared memory tiling.
Actually, if we use a tile-based approach: load a tile of x into shared memory, then iterate over tiles of W, compute partial dots, apply sigmoid, and accumulate to a per-thread local sum. But we need to apply sigmoid after the full dot product is computed.
Hmm, this is tricky. We need the full dot product before applying sigmoid. So we can't sigmoid partial sums.
Let me think differently. For each batch element b:
- We need to compute N intermediate values: y[i] = sigmoid(x[b] · W[i] + bias[i])
- Then sum them up.
With K=N=32768 and M=128.
A custom approach: use one thread block per batch element. Each thread block has T threads. Each thread computes N/T of the dot products. But each dot product is length K=32768. So each thread does (N/T) * K multiply-adds.
Total FLOPs per batch element: N * K = 32768 * 32768 = ~1B.
Total per thread if T=1024: ~1M FMAs. With 128 SMs, we can run ~128 blocks concurrently. Each block doing 1024 threads * 1M ops = 1B ops.
At 67 TFLOPS, 128B FLOPs would take ~1.9ms if we were compute bound and perfectly utilized. But memory bandwidth for loading W: W is 32768 * 32768 * 4 bytes = 4GB. We need to read all of W once (shared across all batch elements actually, since all batch elements use the same W). Plus read x: 128 * 32768 * 4 bytes = 16MB. And read bias: 128KB. Write output: 128 * 4 bytes = 512 bytes.
Total reads ≈ 4GB. At 3.35 TB/s, that's ~1.2ms just for memory. But if we can keep W cached in L2 (which is huge on H100), it might be faster. Actually for 128 batch elements, we read W 128 times? No, if each batch element's block reads W independently, we read W 128 times = 512GB. That would be terrible!
But in reality, on H100 with large L2 cache, if multiple blocks process the same W rows, the cache might help. But 512GB of reads is still a lot for L2 to handle (L2 bandwidth is higher than DRAM but still limited).
Wait - in cuBLAS GEMM, W is read efficiently - each W element is read once and reused across all M batch elements. cuBLAS is highly optimized for exactly this pattern.
So for our custom kernel to beat cuBLAS, it would need to match or exceed cuBLAS's GEMM performance, which is extremely hard.
Alternative approach: use PyTorch's linear which calls cuBLAS, then fuse the sigmoid + reduction ourselves. But as I calculated, the matmul dominates.
Hmm, but let me reconsider. The task says "replace the forward computation" and "run faster than the PyTorch reference". The PyTorch reference does:
1. linear (calls cuBLAS GEMM + bias)
2. sigmoid (elementwise)
3. sum (reduction)
Could there be an advantage to using Tensor Cores for this? The reference linear() in fp32 may not use tensor cores if the PyTorch default is cublas without tensor cores on fp32. But we could use torch.matmul with some flags or use a custom kernel that leverages tensor cores for the exact problem.
Wait, actually in PyTorch, `torch.backends.cuda.matmul.allow_tf32` is True by default since some version. So cuBLAS may already be using TF32 tensor cores for this matmul. The theoretical time at 989 TFLOPS: 275B / 989T = 0.28ms. Plus some overhead.
But the tolerance for fp32 is atol=rtol=1e-4. TF32 has ~1e-3 precision at best (10 mantissa bits vs 24 in fp32), so using TF32 might not pass the 1e-4 tolerance.
Actually, let me reconsider. Even if allow_tf32 is True, PyTorch might not use it for all shapes, or the accumulation might be in fp32, reducing error. But regardless, writing a kernel that uses tensor cores with fp32 accumulation might help.
Actually, I think the best approach is:
1. Recognize that M=128 is small
2. Use a tile-based custom kernel that computes everything in one go
3. The key insight: we can split N (32768) across threads, and have each thread compute partial dot products for a subset of output elements.
But more importantly, we can have one block per batch element, and within each block, compute all N outputs using shared memory for x.
Let me think about the per-block approach:
- Block b processes batch element x[b]
- Load x[b] into shared memory (32768 floats = 128KB). That's within shared memory limit (227KB max per block, 228KB per SM). But we need to be careful - 32768 * 4 = 131072 bytes = 128KB. We have 227KB available per block.
- Then iterate over N in chunks: each thread handles a few outputs. For each output i, we need to compute dot(x[b], W[i]). If W[i] is in global memory, each thread reads K floats.
Wait, but loading x[b] into shared memory allows all threads in the block to share x[b]. Then they iterate over W. But W has N*K elements = 4GB. Reading this from global memory by all 128 blocks is still 512GB of reads unless cached by L2.
With 50MB L2 cache on H100 (I think it's around that), 4GB of W won't fit. So W will be read from DRAM 128 times. That's 512GB. At 3.35 TB/s, that's ~153ms. WAY too slow.
So we absolutely must NOT have each batch element read all of W independently. Instead, we need to batch the processing: have all 128 batch elements cooperate on reading W.
Better approach:
- Divide W into tiles of size TILE_K x TILE_N or TILE_M x TILE_N
- Actually, M=128, N=32768, K=32768
Consider tiles: process W in chunks of, say, 64x1024 or similar.
For each tile of W, load it into shared memory. Then have all batch elements compute their partial contribution using this tile.
Hmm, but M=128 is small. We have 128 SMs. Each SM can handle one tile at a time.
Actually, a more cuBLAS-like approach:
- Use a 2D grid: blocks along N dimension
- Each block computes a chunk of outputs (say N0 to N1) for all batch elements
- Load a tile of W (say 64x128) into shared memory, then all batch elements use it
But implementing a full GEMM kernel is extremely complex. Let me check if there's a simpler win.
Wait, the reference calls `self.linear(x)` which creates a tensor of size (128, 32768), then `torch.sigmoid` on that, then `torch.sum` on that.
What if we just keep the linear in PyTorch (cuBLAS) but fuse sigmoid and sum? The matmul is the bottleneck. Let me think if there's any inefficiency in the reference.
Actually, in PyTorch:
1. `self.linear(x)` does the matmul + bias addition
2. `torch.sigmoid(x)` reads and writes (128, 32768)
3. `torch.sum(x, dim=1, keepdim=True)` reads (128, 32768) and writes (128, 1)
The memory cost of sigmoid + sum is negligible compared to matmul compute. So fusing them won't produce meaningful speedup over cuBLAS GEMM.
Unless... the PyTorch linear + sigmoid + sum are not perfectly overlapping or there's some overread. But for a compute-bound GEMM, this shouldn't matter.
Wait, I should re-check: could we use CUTLASS or write our own kernel that leverages the fact the final output is tiny (128 values), and compute the matmul in a way that reduces memory traffic to the intermediate results?
Actually, I think the biggest win might be possible if PyTorch's `nn.Linear` is using CUTLASS or cuBLAS without tensor cores, or with lower performance. But on H100, cuBLAS is extremely optimized.
Let me think about using Warp Matrix Multiply Accumulate (WMMA) or mma.sync instructions ourselves. But writing a full tensor core kernel in raw CUDA for GEMM is very complex.
Another thought: use `torch.nn.functional.sigmoid_focal_loss` type of fusion? No.
Actually, could we transpose the problem? The weight is (hidden_size, input_size) in `nn.Linear`. The forward is x @ W^T.
Wait, maybe we can use CUTLASS via `torch.utils.cpp_extension`? No, we can't easily include CUTLASS.
Hmm, what if the speedup comes from a custom reduction kernel after a PyTorch matmul? No, reduction is tiny.
Let me reconsider: what if I write a kernel that computes each output element (batch_size, 1) by having each thread block compute one output, and within the block, use warp-level primitives or shared memory for the reduction across N.
Wait! I have an idea. What if we use a simple but highly optimized kernel where:
For each batch element b, we want output[b] = sum_i sigmoid(x[b] · W[i] + bias[i])
We can process this by:
- Each block handles one batch element
- The block has many threads (say 1024)
- Each thread computes the sum for a chunk of i values (e.g., 32 values per thread)
- For each i, we need to compute dot(x, W[i]) which is K=32768 FMAs
If T threads compute for chunk size C = N/T = 32 outputs per thread, then each thread does 32 * 32768 = ~1M FMAs. With 128 blocks on 128 SMs, each doing 1024 threads * 1M FMAs over ~4ms (at 67 TFLOPS). But memory: each block reads all of W for its batch element: 4GB per block * 128 blocks. That's 512GB of reads.
However, if the L2 cache has high bandwidth and the blocks are scheduled such that they overlap on reading the same W elements... but with 128 SMs, each running one block concurrently, and each block reading different parts of W (since threads in a block handle different i values), they might all read different 4GB. Unless W is cached block-wise.
Actually, all blocks read the ENTIRE W, but for different i values so each block reads all of W. But since all blocks read W simultaneously, and each block needs all of W (because each thread in the block handles different i values but all i values for a block span the entire W)... wait no, threads in a block handle different i values, but all of them need all of x[b] and all of W[i] for their respective i values. So each thread only reads a subset of W (the rows it handles). If T=1024 and each thread handles 32 rows, each thread reads 32 * 32768 floats = 4MB. The whole block reads 4GB.
But here's the thing: block 0 reads W[0:32768] to compute for batch element 0. Block 1 reads W[0:32768] to compute for batch element 1. Both blocks read the same W. If they're running concurrently, L2 might help, but realistically not all blocks can fit W in L2.
Wait, H100 L2 cache is 50MB. W is 4GB. So only a tiny fraction of W can be in L2. Most reads will be from HBM. So each block reading W independently means 512GB of reads. That's terrible.
This confirms that we cannot beat cuBLAS for the matmul itself with a naive custom kernel. We must either:
1. Use cuBLAS for matmul and fuse the rest
2. Write a highly optimized GEMM that matches cuBLAS
But if using cuBLAS is required, and the problem is compute-bound, then our only path to speedup is if the reference is somehow inefficient.
Let me re-examine the reference implementation and consider what PyTorch actually does:
1. `self.linear(x)` : calls `addmm` or `mm + add`. For fp32, this calls cuBLAS's `cublasSgemm`.
2. `torch.sigmoid(x)` : elementwise kernel
3. `torch.sum(x, dim=1, keepdim=True)` : reduction kernel
The sigmoid writes 128*32768*4 = 16MB. Sum reads 16MB and writes 0.5KB. These are essentially free compared to the 275B FLOPs matmul.
BUT - could we avoid the memory traffic for the intermediate (128, 32768) tensor? If we fuse sigmoid+sum in a custom kernel after the matmul, we still have to write the matmul output from cuBLAS and read it in our kernel. That's still 16MB each way.
What if we use a custom CUTLASS-like epilogue? We can't easily use CUTLASS from load_inline.
Wait, maybe there's a simpler insight. On H100, cuBLAS uses TF32 by default when `torch.backends.cuda.matmul.allow_tf32` is True. But this compute is very heavy and might take ~0.3ms. The actual reference time might be dominated by matmul kernel launch or some inefficiency.
What if the actual reference doesn't use cuBLAS optimally because of the bias addition? PyTorch linear adds bias with a separate kernel in some code paths.
Actually, in PyTorch, `nn.Linear` might be implemented as `F.linear`, which does `input @ weight.t() + bias`. When input is 2D, `input @ weight.t()` calls `addmm` with bias. Wait, actually PyTorch's `F.linear(input, weight, bias)` is equivalent to `input.matmul(weight.t())` and then add bias if bias is not None. Let me check...
Actually, PyTorch's `nn.Linear.forward()` typically calls `F.linear(input, self.weight, self.bias)`, which is:
```python
output = input.matmul(weight.t())
if bias is not None:
output += bias
```
For 2D input, matmul calls into cuBLAS. But sometimes PyTorch might dispatch to a more generic path. Still, for these sizes, it should use cuBLAS GEMM.
Now the bias addition is elementwise and adds memory traffic. But cuBLAS can sometimes fuse bias via `C = alpha * A * B + beta * C` if bias is already in C, but here we need `C = A * B + bias`. PyTorch might do a separate bias kernel or use `addmm` which is: `addmm(input, mat1, mat2, beta=1, alpha=1)` = `beta * input + alpha * (mat1 @ mat2)`.
Actually, `torch.addmm(bias, input, weight.t())` where bias is expanded to (M, N) or treated as 1D... this might be a very fast fused operation with cuBLAS!
Wait, `F.linear` source code is: `input @ weight.t() + bias`. For 2D input:
```python
return torch.addmm(bias, input, weight.t())
```
Yes! PyTorch uses `addmm` which is "add matrix product" and cuBLAS can handle this efficiently (often fusing the bias addition into the GEMM epilogue or doing it as part of the main kernel). So this is already very fast.
OK so the linear part is cuBLAS GEMM with bias. Then sigmoid. Then sum.
The only way to significantly beat this is to fuse all three operations so that we never write the full (128, 32768) tensor to global memory. That is, fuse matmul -> sigmoid -> sum into one kernel that writes only 128 floats.
But can we write a kernel that does this efficiently? As discussed, a naive approach would have terrible memory reuse for W. We need a tiling strategy.
Wait, what about using flash-attention-like techniques? No, that's for attention.
Let me think about a tiled approach more carefully.
We have:
- x: (M=128, K=32768)
- W: (N=32768, K=32768) - stored as (K, N) or transposed to (N, K)
- For each m in [0, M): output[m] = sum_{n=0}^{N-1} sigmoid(sum_{k=0}^{K-1} x[m,k] * W[n,k] + bias[n])
If W is stored as (N, K) in row-major, then W[n] is a contiguous row of length K.
For each output element m, we need to compute N dot products.
Tiling approach for the entire computation:
- Divide K into tiles of size BK (e.g., 64 or 128)
- Divide N into tiles of size BN (e.g., 64)
- Process the output by computing partial sums for tiles of N using tiles of K
Actually a better parallelization:
- Grid: (M, N/BN) blocks. Each block computes a partial sum for one (m, n_tile) pair.
- But we need the FULL dot product before sigmoid, so we can't do block-level partial sigmoids.
Hmm, that means each sigmoid needs the complete dot product, so we need one thread/cooperative group to compute the entire dot product for one n before applying sigmoid.
Alternative: process the problem as a batched 1×K by K×N reduction-after-sigmoid.
Since M=128 is small, let's parallelize over (M, N). There are 128 * 32768 = ~4.2M output elements before the final sum. For each (m, n):
- y[m,n] = sigmoid(x[m] · W[n] + bias[n])
To compute one y[m,n], we need a dot product of length 32768.
If we use one thread per y[m,n], each thread does 32768 FMAs + 1 sigmoid. Total FMAs: 4.2M * 32768 ≈ 137B FMAs = 275B FLOPs. Same as GEMM.
At 67 TFLOPS, this takes ~4.1ms. But can we feed threads fast enough?
Each thread reads K=32768 floats from x and K=32768 floats from W[n]. But x[m] is shared among all threads computing for the same m. If we have T threads computing for the same m but different n values, they all need to read x[m] (cacheable) and their respective W[n] (different).
If we process, say, 128 SMs each with 1024 threads, total 131072 threads. Each thread handles one (m,n). There are 4.2M (m,n) pairs. So each thread would handle ~32 (m,n) pairs.
Wait, that's actually not a bad approach! Let me think again.
Total work: M*N*K = 128 * 32768 * 32768 FMAs.
With 131072 active threads (all SMs fully utilized), each thread does:
128 * 32768 * 32768 / 131072 = 128 * 32768 * 0.25 = 1,048,576 FMAs.
That's about 1M FMAs per thread. At ~2 FLOPs/cycle/core, with clock ~2GHz, each thread would take ~0.5M cycles if perfectly pipelined. That's 0.25ms per thread. With many threads, this could work.
But memory bandwidth: each thread reads one contiguous row of W: 32768 * 4 = 128KB per thread. There are 131072 threads, so total W reads: 131072 * 128KB = 16GB.
Wait, that's much better than 512GB! Because each thread handles one (m,n) and only reads W[n], and there are 131072 threads handling 131072 different n values. But W has N=32768 rows, so with 131072 threads, each W[n] is read by exactly M=128 threads (one per batch element). Total W reads: 32768 rows * 128 reads * 128KB = 512GB. Same as before!
Hmm yes, because there are M threads reading each W[n]. Total reads of W: M * N * K floats * 4 bytes = 128 * 4GB = 512GB. Plus x reads: M*K floats read by all threads working on that m, but cached.
So total memory: ~512GB. At 3.35 TB/s = 153ms. That's the fundamental issue: we need to read W M=128 times.
cuBLAS GEMM solves this by having each block of threads compute a tile of the output, reusing W across multiple M values. In cuBLAS, W is typically loaded into shared memory and reused for multiple output rows.
So to match cuBLAS, our custom kernel also needs to tile W and reuse it across M batch elements.
Here's a good tiling strategy:
- One block computes a tile of size TM × TN output elements
- Load a tile of x of size TM × BK into shared memory
- Load a tile of W of size TN × BK into shared memory (or transposed: BK × TN)
- Compute partial dot products
- Iterate over K in tiles
For our case, since M=128 is small, we can set TM=M=128 so one block handles all batch elements for a subset of N outputs. Then we iterate over K.
Dimentions:
- TM = 128 (all batch elements)
- We need to choose TN and BK such that tiles fit in shared memory.
Shared memory needed:
- x tile: TM × BK floats = 128 * BK * 4 bytes
- W tile: TN × BK floats = TN * BK * 4 bytes
- Total: 4 * BK * (128 + TN) bytes
If BK=64, TN=64: total = 4 * 64 * 192 = 49152 bytes = 48KB ✓
If BK=64, TN=128: total = 4 * 64 * 256 = 65536 bytes = 64KB ✓
If BK=128, TN=128: total = 4 * 128 * 256 = 131072 bytes = 128KB ✓
Then for each output tile (TN outputs), we:
1. Load W[n_start:n_start+TN, k_start:k_start+BK] into shared memory
2. For each m in [0, M): load x[m, k_start:k_start+BK] into shared memory (actually all M=128 rows need to be loaded)
3. Compute partial dot products for all M batch elements and TN outputs
4. Iterate over K
Since M=128 and TN=128, each block computes 128 * 128 output values. There are N/TN = 256 such blocks. These 256 blocks can run in parallel across 128 SMs, with 2 waves.
This is essentially a standard GEMM kernel! Then after accumulating the dot product for each (m,n), we apply sigmoid and do the sum over n.
For the sum over n, we can either:
1. Write all y[m,n] to global memory, then do a reduction (which would still need a separate kernel)
2. Keep the sum reduction in registers/shared memory within each block, but since one block only handles TN outputs, we'd need another reduction across blocks.
Actually, since a block only handles a slice of N, we can have each block compute partial sums for its TN outputs and write them to global memory. Then launch another kernel or use atomics to sum across the N dimension. But wait, we can do better: have the blocks write their partial sums for the M batch elements to global memory (M * N_blocks floats = 128 * 256 = 32768 floats = 128KB). Then do a quick reduce kernel. But writing 128KB is essentially free.
Wait, we don't even need a second kernel if we use cooperative groups or grid sync... but grid sync is complex.
Actually, we can compute the final sum in two steps:
Step 1: Tiled GEMM-like kernel that computes dot products and applies sigmoid, writes a partial output of shape (M, N/TN) or (M, N) - but writing full (M, N) would cost 16MB which might still be acceptable (and we'd be no better than PyTorch reference).
Wait, no. Writing and reading 16MB costs ~5 microseconds at 3.35 TB/s. For a 4ms kernel, that's negligible. But our goal is to beat cuBLAS GEMM + tiny kernels, which take ~4ms total. If our GEMM kernel is as fast as cuBLAS but writes 16MB, and then we have another kernel, we're at best equal, probably slower because our GEMM won't be as good as cuBLAS.
Wait, so if we can't beat cuBLAS GEMM, what can we do?
Actually, maybe I should just benchmark first by writing a kernel and seeing. Let me also consider:
- Could PyTorch's reference be surprisingly slow because it doesn't use cuBLAS optimally?
- Could there be significant overhead from the Python side or memory allocations?
Actually, I notice that in the reference, `self.linear` stores weight as (hidden_size, input_size) = (32768, 32768). When we call `self.linear(x)`, PyTorch computes `x @ weight.t()`. Since weight is stored as (out_features, in_features), the transpose is handled. cuBLAS supports this natively by using the transpose flag.
Let me try writing a relatively simple fused kernel and see how it compares.
Let me think about a practical kernel design that maximizes what we can do:
Actually, for M=128, K=32768, N=32768, this is very close to square GEMM. cuBLAS should handle this near perfectly. And anything I write in raw CUDA won't come close to cuBLAS.
UNLESS - I use Tensor Cores via the mma PTX instructions or CUDA WMMA API. But that's extremely complex for a load_inline source.
Wait, actually there's another path. What if PyTorch's reference is using the Transformer Engine or something that makes it slower? No, this is a simple linear layer.
Let me reconsider whether the problem is compute bound at M=128.
FLOPs = 2 * 128 * 32768 * 32768 = 274.9e9
At H100 FP32 67 TFLOPS, that's 4.1ms.
At H100 TF32 tensor core 989 TFLOPS, that's 0.278ms.
The question is: does PyTorch use tensor cores for fp32 linear? On H100, `torch.backends.cuda.matmul.allow_tf32 = True` by default on newer PyTorch. If TF32 is used, the matmul takes ~0.3ms and our task is much harder. If pure FP32 is used (~4ms), maybe we can fuse operations to save ~10 microseconds, which is a 0.2% speedup. That probably won't "run faster" in a measurable way.
Unless... we write a kernel that uses tensor cores for FP32 accumulation but in a way that still passes 1e-4 tolerance. With TP32 (19-bit precision) the error might be around 1e-3, which could fail 1e-4 tolerance.
Actually, wait. Let me reconsider: M=128 is small. Maybe PyTorch's GEMM launch overhead or workspace allocation is significant. For a 0.3ms TF32 kernel, 0.05ms overhead is 15%. For a 4ms FP32 kernel, it's 1.25%.
Hmm. What if we use CUTLASS-like epilogue fusion via PyTorch's own APIs? PyTorch doesn't expose this easily.
Another idea: instead of computing the full GEMM, notice that the final output is very small (128 values). If we could split the work across the M dimension and reduce... no, we need the full N dimension per M.
OK, let me think about this completely differently. What if the trick is that `torch.sigmoid(x) + sum` can be computed in a way that avoids creating large intermediates, but the main speedup comes from something else?
Actually, I just realized something. In the reference model:
```python
x = self.linear(x) # (128, 32768)
x = torch.sigmoid(x) # (128, 32768)
x = torch.sum(x, dim=1, keepdim=True) # (128, 1)
```
The sum is over dim=1 (the hidden_size dimension). That's a simple reduction.
Could there be a win from using `torch.cuda.amp.autocast()` or something? No, we can't change precision.
What about using a custom kernel that uses float16/bfloat16 internally? The spec says "produce numerically equivalent outputs at fp32 precision", and tolerance is 1e-4 for fp32. If we use tensor cores with FP16 and FP32 accumulation, the accuracy might be around 1e-3 which would fail. But if we use a block-wise mixed precision with iterative refinement that compensates... no, that's too complex.
Wait, what if the secret is that we can compute this without actually doing M*N*K FLOPs?
output[b] = sum_i sigmoid(x[b] · W[i] + bias[i])
= sum_i sigmoid(sum_j x[b,j] * W[i,j] + bias[i])
Is there any mathematical simplification? No, sigmoid is nonlinear and doesn't distribute over sum, and sum over sigmoid doesn't simplify.
What about using a Lookup Table for sigmoid? For 32768 elements per row, computing `1/(1+exp(-x))` 32768 * 128 = 4.2M times. exp() is expensive but 4.2M times is not that much on GPU.
Let me step back. Since I can't easily beat cuBLAS GEMM, let's see if I can write a fused kernel that:
1. For each batch element, uses shared memory for x
2. Processes W in tiles to ensure W is read only once (or close to once)
3. Computes sigmoid and accumulates sum on-the-fly
The key insight for reading W only once: all batch elements share the same W. So we process W in tiles, and for each tile, all batch elements use it.
Actually, the standard GEMM already does this! Each W tile is loaded once and reused for multiple M values. The difference from my naive per-row approach is that cuBLAS tiles the computation.
So I need to write a tiled GEMM kernel where:
- W is loaded in tiles and reused across all M=128 batch elements
- The output is accumulated per (m,n)
- After the full dot product is computed, sigmoid is applied and accumulated into the final sum per m
For the final sum: if each block handles some n-range, we can accumulate into a per-block shared memory buffer of size M, and then atomically add to the final output of size M. Or each block writes its M partial sums, then another kernel reduces.
Actually, here's a nice approach:
- Grid has blocks along the N dimension. Each block handles a chunk of N (say 256 or 512 outputs).
- Each block has threads organized to compute a tile of (M_chunk, N_chunk) outputs.
- X is loaded into shared memory in tiles
- W is loaded into shared memory in tiles
- After computing all dot products for the (M_chunk, N_chunk) tile:
- Apply sigmoid to each accumulated value
- Each thread locally sums its outputs along the N dimension
- Then within the block, reduce these partial sums to one sum per M value
- Atomically add to output[m]
This way, we never write the full (M, N) matrix to global memory! We only write M=128 floats.
Now the key challenge is getting GEMM performance close to cuBLAS for M=128. Since M is small, but K and N are large, we need to maximize throughput.
For a standard GEMM where C = A @ B:
- A: (M, K) = (128, 32768)
- B: (K, N) = (32768, 32768) - but wait, W in PyTorch linear is (N, K) = (32768, 32768)
Actually, `x @ weight.t()`:
- x is (M, K) where M=128, K=32768
- weight is (N, K) where N=32768, K=32768
- weight.t() is (K, N)
- x @ weight.t() is (M, N) = (128, 32768)
In terms of the standard GEMM C = A @ B:
- A = x: M × K = 128 × 32768
- B = weight.t(): K × N = 32768 × 32768
- C = output: M × N = 128 × 32768
For M=128, N=32768, K=32768.
A standard CUDA GEMM kernel for large matrices uses blocks that compute tiles of C. Since M=128 is relatively small, we might use a block size like:
- BM = 8 or 16 or 32 (batch elements per block)
- BN = 128 or 256 (output features per block)
- BK = 16 or 32 or 64 (inner dimension tile)
Then the grid is (M/BM, N/BN).
Shared memory:
- A tile: BM × BK
- B tile: BK × BN (but B is K×N, so we load BK rows and BN columns, or transpose)
Actually, since B = weight.t() and weight is stored as (N, K), B[k, n] = weight[n, k]. So B is column-major with respect to weight. This might hurt cache locality. But cuBLAS handles this.
Actually, PyTorch stores weight as (out_features, in_features). For linear, we compute x @ weight.t(). In GEMM terms with column-major (cuBLAS default), let me think...
Actually for simplicity, let me think in terms of row-major where:
- x[m, k] is row-major: x[m*K + k]
- weight[n, k] is row-major: weight[n*K + k]
- Output y[m, n] = sum_k x[m, k] * weight[n, k] + bias[n]
So it's like computing dot products between each row of x and each row of weight, then adding bias. Wait, that's exactly it: y[m, n] = dot(x[m], weight[n]) + bias[n].
For each n, we need to dot with all m.
So we can think of it as: for each weight row n, compute its dot product with all x rows. Since M=128 is small, we can load 128 rows of x? No, x only has 128 rows total. We can load all x into shared memory: 128 * 32768 * 4 bytes = 16MB. That's way too big (shared memory limit is 227KB per block).
We can load x in tiles along K. For each tile of K, load a slice of x (M × BK) and a slice of weight (some rows × BK).
OK let me try a specific tiling strategy for a custom kernel:
**Kernel Design:**
- Each block computes outputs for all M=128 batch elements and a chunk of N outputs.
- Since M=128 is small, we might have each block handle a chunk of N outputs, say BN=64 or 128.
- Within a block, each thread computes outputs for some subset of (m, n).
- We tile over K with BK=64.
Block dimensions:
- BlockDim.x = 128 threads
- BlockDim.y = 1 or 2
Per block:
- Load x_tile of size M × BK = 128 × 64 floats = 32768 bytes = 32KB into shared memory
- Load w_tile of size BN × BK = 128 × 64 floats = 32768 bytes = 32KB into shared memory
- Total shared memory: 64KB ✓
Wait, but for BN=128, w_tile is 128*64*4 = 32KB. Total = 64KB. Good.
Number of blocks: N / BN = 32768 / 128 = 256 blocks.
Grid: (256, 1)
Each block:
1. Initialize accumulators for M × BN output values (128 × 128 = 16384 floats). That's too many for registers! If each thread handles multiple outputs, say 16 threads per block, each thread needs 128*128/16 = 1024 registers. That's too many.
Wait, if we have 128 threads in a block, and each thread computes all M outputs for one n-value... Then each thread handles 1 n-value and all 128 m-values. Accumulators: 128 floats per thread. At 4 bytes each, 512 bytes per thread. With 128 threads, total block registers: 128 * 512 / 4 = 16384 registers. That's 16K 4-byte registers. Each SM has 64K registers, so this should fit.
Wait, each thread computes y[m, n] for all m and one fixed n (or a few n values). The dot product for each y[m,n] is accumulated independently. So each thread needs M accumulators = 128 floats = 512 bytes. 128 threads means the block uses 128 * 128 registers = 16384 registers. That fits in 64K! Plus some for temporaries, index calculations, etc.
Actually, H100 has 64K 32-bit registers per SM, and up to 255 per thread. A thread with 128 accumulators + temporaries might use 150 registers, which is well under 255.
So here's the design:
- Grid: (N/BN, 1) where BN = number of n-values per block
- Block: BN threads (or some multiple)
- Each block handles BN consecutive n-values
- Each thread handles one n-value and all M=128 m-values
- For a fixed n, thread loads weight[n, :] and computes dot products with all x[m, :]
Wait, but if each thread handles one n and all m, then for a given K tile [k_start, k_start+BK), the thread needs weight[n, k_start:k_start+BK]. And it needs x[m, k_start:k_start+BK] for all m.
So shared memory approach:
- Load x_tile = x[:, k_start:k_start+BK] into shared memory: M × BK floats. All threads in the block need this.
- Load w_tile = weight[n_start:n_start+BN, k_start:k_start+BK] into shared memory: BN × BK floats. Each thread only needs one row.
- Then each thread computes partial dot products for its n and all m using x_tile and its row from w_tile.
This is elegant!
Let me work out the numbers:
- M = 128, K = 32768, N = 32768
- BK = 64 (tile size along K)
- BN = 128 (number of n-values per block)
- Block size = BN = 128 threads
Per block:
- Shared mem for x_tile: 128 × 64 × 4 = 32768 bytes
- Shared mem for w_tile: 128 × 64 × 4 = 32768 bytes
- Total shared mem: 65536 bytes = 64KB ✓
Each block:
- Number of K iterations: K / BK = 32768 / 64 = 512
- Per K iteration, each thread computes M=128 partial dot products
- Total FMAs per block: BN × M × K = 128 × 128 × 32768 = ~537M
- Total FMAs overall: N × M × K = 32768 × 128 × 32768 = ~137B (matches!)
Each thread:
- Number of iterations: K / BK = 512
- Per iteration: compute M dot products, each of BK elements
- But x_tile is in shared memory, and the thread gets one element of w_tile at a time
- For each k in BK and each m in M: accumulator[m] += x_tile[m, k] * w_tile[thread_n, k]
Wait, the inner loop would be:
```cpp
for (int k = 0; k < BK; k++) {
float w_val = w_tile[threadIdx.x * BK + k];
for (int m = 0; m < M; m++) {
accum[m] += x_tile[m * BK + k] * w_val;
}
}
```
This accesses x_tile[m * BK + k] for all m with the same k. In shared memory, x_tile is stored with m varying fastest. So x_tile[m * BK + k] for different m at fixed k are strided by BK. If BK=64, these accesses are coalesced for the threads? Actually within a single thread, we're just reading sequentially through m.
Wait, but all threads read w_tile at the same time (different rows), which is coalesced. x_tile is read by all threads for each m. The same x_tile element is read by all threads simultaneously. That's a broadcast, which is efficient in shared memory.
Yes! This should work well.
After all K tiles are processed, each thread has M accumulators (one per batch element) for its n value.
Then: for each m, apply sigmoid(accum[m] + bias[n]) and add to a per-thread sum.
Then reduce the per-thread sums within the block to get per-m partial sums.
Finally, atomically add these partial sums to the global output.
Let me refine:
```cpp
// Per thread accumulators for M=128 batch elements
float accum[128]; // 128 floats per thread
// For each K tile
for (int k_tile = 0; k_tile < K_tiles; k_tile++) {
// Load x_tile into shared memory: M × BK
// Load w_tile into shared memory: BN × BK
__syncthreads();
// Compute partial dots
for (int k = 0; k < BK; k++) {
float w_val = w_tile[threadIdx.x * BK + k];
#pragma unroll
for (int m = 0; m < M; m++) {
accum[m] += x_tile[m * BK + k] * w_val;
}
}
__syncthreads();
}
// Now apply sigmoid and accumulate into block-local sums
float block_sum[M] = {0}; // in shared memory, size M floats
float thread_sum[M] = {0};
for (int m = 0; m < M; m++) {
float val = accum[m] + bias[n];
thread_sum[m] = 1.0f / (1.0f + expf(-val)); // sigmoid
}
// Reduce thread_sum across threads in block
// We need to sum thread_sum[m] for each m across all BN threads
// Store in shared memory
__shared__ float block_sums[M]; // M = 128
// For each m, all threads contribute
for (int m = 0; m < M; m++) {
float sum = thread_sum[m];
// Warp shuffle reduce
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
}
if (threadIdx.x % warpSize == 0) {
block_sums[m] += sum; // Wait, this is wrong for multiple warps
}
// Actually need inter-warp reduction too
}
```
Hmm, the reduction is a bit tricky because we have BN=128 threads = 4 warps of 32 threads. We need to reduce across all 128 threads.
Better approach for reduction:
- Each warp reduces its 32 values using warp shuffle
- The first thread of each warp writes to a shared memory array warp_sums[4]
- After all warps done, one warp reads warp_sums and reduces
- Result is the block sum for that m
- Then atomically add to output[m]
But we're doing this for M=128 different values, so we'd have 128 parallel reductions.
Actually, we can transpose the reduction:
- For each m, all BN threads have a partial value
- We could also just write all BN * M partial sums and reduce in a second kernel. BN=128, M=128, so 16384 floats = 64KB. Very small!
Or even simpler: have each block write its M partial sums to global memory at an offset `blockIdx.x * M`. Then launch a tiny reduction kernel that sums across blocks.
But launching a second kernel adds overhead. For M=128 and 256 blocks, the reduction is trivial. We can even avoid a second kernel by using `atomicAdd` for the M output values.
```cpp
// After warp reductions, block_sums[m] contains the sum for this block
for (int m = 0; m < M; m++) {
if (threadIdx.x == 0) {
atomicAdd(&output[m], block_sums[m]);
}
}
```
Only 128 atomic adds per block, and 256 blocks. That's fine.
Wait, but atomics from 256 blocks to the same 128 locations... that's 256*128 = 32768 atomic operations. On H100, atomics to global memory are reasonably fast, and since it's contention on only 128 addresses... actually it might be contentious. But we can use `atomicAdd` with `__half` or we could do:
Alternative: write partial_sums[blockIdx.x * M + m] = block_sum[m]. Then after the kernel, sum on CPU or with a tiny kernel. But the sum is only 256*128 = 32768 floats. Doing this on CPU takes almost no time. Or if we want to stay on GPU, a single block can handle it.
Actually, since we're in `ModelNew.forward()`, we can just call one custom kernel that writes partial_sums, then a second tiny kernel or use PyTorch to sum. But launching two kernels adds overhead. The atomic approach avoids the second kernel.
Let me reconsider using atomics: 256 blocks each adding to 128 output locations. If all blocks finish at different times, contention is low. With 128 SMs, at most 128 blocks run concurrently, each adding to 128 locations. The memory addresses are different for different m. So there's no contention across different m values. Each m gets at most 128 concurrent atomic adds. On H100, this should be fine.
OK, let me code up this kernel.
One concern: M=128 hardcoded loops. We should probably make it a compile-time constant with templates, or just `#define M 128`. But since the input batch size comes at runtime, we might pass it as a parameter. However with M=128 from the problem statement, we can use it as a compile-time constant for unrolling.
But actually, let's make the kernel work for any M (though we'll specialize for M=128). Since we know M=128, we can use `#define M 128`.
Wait, is M guaranteed to be 128? The problem says:
```python
batch_size = 128
input_size = 32768
hidden_size = 32768
```
And `get_inputs()` always returns `torch.rand(batch_size, input_size)`. So yes, M is always 128.
Similarly, `input_size` and `hidden_size` are fixed at 32768. So K=N=32768 always. This simplifies things!
We can completely hardcode the dimensions. This is a huge advantage because:
- We can fully unroll loops
- We can optimize shared memory sizes exactly
- We don't need generic indexing
Let me write a highly optimized kernel with hardcoded sizes:
```cuda
#define M 128
#define K 32768
#define N 32768
#define BK 64
#define BN 128
#define NUM_K_TILES (K / BK) // 512
#define NUM_BLOCKS (N / BN) // 256
```
Kernel:
```cuda
__global__ void fused_linear_sigmoid_sum(
const float* __restrict__ x, // (M, K)
const float* __restrict__ weight, // (N, K) - note: stored as row-major, each row is a weight vector
const float* __restrict__ bias, // (N,)
float* __restrict__ output // (M,)
) {
__shared__ float x_tile[M * BK];
__shared__ float w_tile[BN * BK];
int n_start = blockIdx.x * BN;
int tid = threadIdx.x;
// Each thread handles one n value: n = n_start + tid
int n = n_start + tid;
// Accumulators for M batch elements
float accum[M];
#pragma unroll
for (int m = 0; m < M; m++) accum[m] = 0.0f;
// Iterate over K tiles
for (int kt = 0; kt < NUM_K_TILES; kt++) {
int k_off = kt * BK;
// Load x_tile: M rows × BK cols
// x is (M, K), stored row-major
// We need x[m, k_off + k] for m in [0,M), k in [0,BK)
// Each thread can load multiple elements
// Total elements: M * BK = 8192
// With BN = 128 threads, each thread loads 64 elements
for (int idx = tid; idx < M * BK; idx += BN) {
int m = idx / BK;
int k = idx % BK;
x_tile[idx] = x[m * K + k_off + k];
}
// Load w_tile: BN rows × BK cols
// weight is (N, K), row-major
// We need weight[n_start + bn, k_off + k] for bn in [0,BN), k in [0,BK)
// Total elements: BN * BK = 8192
// Each thread loads 64 elements
for (int idx = tid; idx < BN * BK; idx += BN) {
int bn = idx / BK;
int k = idx % BK;
w_tile[idx] = weight[(n_start + bn) * K + k_off + k];
}
__syncthreads();
// Compute partial dot products
// Thread tid computes for n = n_start + tid
#pragma unroll
for (int k = 0; k < BK; k++) {
float w_val = w_tile[tid * BK + k];
#pragma unroll
for (int m = 0; m < M; m++) {
accum[m] += x_tile[m * BK + k] * w_val;
}
}
__syncthreads();
}
// Apply sigmoid and sum within thread
float thread_sum[M];
#pragma unroll
for (int m = 0; m < M; m++) {
float val = accum[m] + bias[n];
thread_sum[m] = 1.0f / (1.0f + expf(-val));
}
// Warp-level reduction for each m
// We have 4 warps (128 threads)
// For each m, reduce across 128 threads
__shared__ float warp_sum[4 * M]; // 4 warps × 128 values
for (int m = 0; m < M; m++) {
float sum = thread_sum[m];
// Warp reduce
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
}
int warp_id = tid / 32;
if (tid % 32 == 0) {
warp_sum[warp_id * M + m] = sum;
}
}
__syncthreads();
// Warp 0 reduces the 4 warp sums
if (tid < 32) {
for (int m = 0; m < M; m++) {
float sum = (tid < 4) ? warp_sum[tid * M + m] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
}
if (tid == 0) {
atomicAdd(&output[m], sum);
}
}
}
}
```
Wait, I need to be careful. `__shfl_down_sync` with offset 16, 8, 4, 2, 1 for warp of 32 threads. Let me write it correctly:
```cpp
for (int offset = 16; offset > 0; offset >>= 1) {
sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
}
```
Actually, warp size is 32. So:
```cpp
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
}
```
Now, a few issues:
1. `atomicAdd` on output[m] - but what about the initial value of output? We need to initialize output to 0.
2. With 128 threads and each loading 64 elements for x_tile and 64 for w_tile, the load loops have 64 iterations each per K tile. With 512 K tiles, that's 32768 iterations for loading, plus compute. The compute is 512 * 64 * 128 = 4,194,304 FMAs per thread.
3. Memory access pattern for loading x_tile: `x[m * K + k_off + k]`. When `k_off` changes, we're accessing different strides. But since x is 16MB total, and we access it sequentially by tiles, it should be fine. But note that `m * K + k_off + k` means for a fixed m and k_off, as k varies, we access contiguous memory. Good.
4. Memory access for w_tile: `weight[(n_start + bn) * K + k_off + k]`. For fixed bn and k_off, as k varies, contiguous. Good.
Wait, but x_tile loading: each thread loads 64 elements. Thread 0 loads x[0*K + k_off + 0] to x[0*K + k_off + BK-1] in the first iteration of its loop? Let me check:
```cpp
for (int idx = tid; idx < M * BK; idx += BN) {
int m = idx / BK;
int k = idx % BK;
x_tile[idx] = x[m * K + k_off + k];
}
```
With tid=0: idx=0, BN=128: next idx=128. m=0/BK=0, k=0. For idx=128: m=128/64=2, k=0. Thread 0 loads every 128th element. That is NOT contiguous in x!
This is actually OK because threads are accessing memory in a strided pattern. With BN=128 threads, thread t accesses x at offset t, t+128, t+256, etc. In row-major x with shape (128, 32768), `idx = tid = 0` maps to m=0, k=0. `idx=128` maps to m=2, k=0. So thread 0 reads x[0,0], x[2,0], x[4,0], ... which are strided by 2*K = 65536 floats = 256KB apart. That's terrible for memory coalescing!
We want coalesced access. Let me redesign the loading.
For x_tile of shape (M, BK) = (128, 64), stored in shared memory as row-major (m * BK + k):
- We want adjacent threads to load adjacent elements from global memory.
- Thread t should load element t (or nearby) from the linear view.
Global memory layout of x: x[m, k] = x[m * K + k_off + k]
If we think of the tile as a contiguous block of M*BK elements laid out in memory... but x is NOT contiguous in the (m, k) region because K=32768 and BK=64. For m=0, elements x[0, k_off] to x[0, k_off+63] are contiguous (64 floats). Then x[1, k_off] is at x[1*K + k_off], which is 32768 floats away.
So x_tile is stored non-contiguously in global memory. To load it efficiently:
- We need threads in a warp to access contiguous or nearby addresses
- But x[0,0:64] and x[1,0:64] are 32768 floats apart
Actually, multiple warps can handle different m values. Let me think differently.
For loading x_tile, we can have each warp load one row. Warp 0 (threads 0-31) loads x[0, k_off:k_off+BK]. Warp 1 loads x[1, k_off:k_off+BK], etc. There are M=128 rows and 32 threads per warp, so 4 warps per row? Or 1 warp per row with each thread loading 2 elements.
Actually, BK=64, so each row has 64 elements. A warp of 32 threads can load 2 elements per thread to get 64 elements. So 1 warp per row, 128 rows = 128 warps. But we only have 128 threads = 4 warps.
So with 4 warps and 128 rows, each warp handles 32 rows. Each row has 64 elements. A warp of 32 threads loads 1 element per thread, so 2 iterations per row, 64 rows total? No, each warp handles 32 rows, 64 elements each = 2048 elements. With 32 threads, each thread loads 64 elements.
Wait, with 128 threads total and 128 rows of 64 elements each = 8192 elements total. Each thread loads 64 elements. With 32 threads per warp, each warp loads 2048 elements = 32 rows.
For coalescing within a warp: if threads 0-31 each load elements from their assigned row, and adjacent threads load adjacent k values, that would work. But what if threads in a warp access different rows? x[m, k] and x[m+1, k] are 32768 floats apart, so threads accessing different m values causes strided access within the warp.
Best approach: assign one warp to one row. But we only have 4 warps and 128 rows. So 32 rows per warp. Threads in a warp process the same row but different k values. So warp 0 processes rows 0-31? No, 32 threads in a warp can only work on one row at a time for perfect coalescing.
Wait, each warp has 32 threads. If all 32 threads load from the same row m but different k (e.g., thread t loads x[m, k_off + t + offset]), then the loads are contiguous in global memory. Perfect coalescing!
So: process x row by row. For each row m, have one warp load all BK elements. With BK=64, the warp loads 64 elements = 2 elements per thread.
To load 128 rows, we need 128 warps. But we only have 4 warps (128 threads). So we can only load 4 rows in parallel. That's fine, it just means the loading takes 128/4 = 32 iterations, with each warp switching to a new row each time.
More precisely:
```cpp
for (int m_iter = 0; m_iter < M; m_iter += num_warps) {
int m = m_iter + warp_id;
if (m < M) {
// Each thread in warp loads 2 elements
for (int i = lane_id; i < BK; i += warpSize) {
x_tile[m * BK + i] = x[m * K + k_off + i];
if (i + warpSize < BK) {
x_tile[m * BK + i + warpSize] = x[m * K + k_off + i + warpSize];
}
}
}
}
```
With warpSize=32 and BK=64, each thread loads 2 elements. Since all threads in the warp load from the same row, and adjacent threads load adjacent k, this is perfectly coalesced.
Similarly for w_tile: weight has N rows, each of length K. For a block handling BN=128 rows starting at n_start, and loading a tile of BK elements per row, each warp can load one row. Since BN=128 and num_warps=4, we need 128/4 = 32 iterations.
```cpp
for (int bn_iter = 0; bn_iter < BN; bn_iter += num_warps) {
int bn = bn_iter + warp_id;
int n = n_start + bn;
if (bn < BN) {
for (int i = lane_id; i < BK; i += warpSize) {
w_tile[bn * BK + i] = weight[n * K + k_off + i];
if (i + warpSize < BK) {
w_tile[bn * BK + i + warpSize] = weight[n * K + k_off + i + warpSize];
}
}
}
}
```
This is perfectly coalesced for w_tile too since adjacent threads load adjacent k values from the same row.
Wait, weight has shape (N, K) = (32768, 32768). Row n starts at n*K. So weight[n, k_off + i] is at n*K + k_off + i. Adjacent threads (lane_id and lane_id+1) access n*K + k_off + i and n*K + k_off + i+1. Perfectly contiguous. Great!
But using warp-specific code with __shfl and lane_id makes the kernel more complex. Let me see if we can simplify.
Actually, with simple flat indexing and careful choice, we might still get decent coalescing. Let me reconsider.
If we use the flat loop:
```cpp
for (int idx = tid; idx < M * BK; idx += BN) {
int m = idx / BK;
int k = idx % BK;
x_tile[idx] = x[m * K + k_off + k];
}
```
With BN=128, thread 0 loads idx=0, 128, 256, 384, ...
These map to (m, k): (0,0), (2,0), (4,0), (6,0), ...
So thread 0 accesses x[0], x[2*K], x[4*K], x[6*K], ...
Thread 1 accesses x[1], x[1*K+1? No], wait:
idx=1: m=0, k=1 -> x[K_off+1]
idx=129: m=2, k=1 -> x[2*K + K_off+1]
So thread 1 accesses x[K_off+1], x[2*K+K_off+1], etc.
Warp 0 (threads 0-31) accesses:
Thread 0: x[0*K + k_off + 0], x[2*K + k_off + 0], x[4*K + k_off + 0], ...
Thread 1: x[0*K + k_off + 1], x[2*K + k_off + 1], x[4*K + k_off + 1], ...
...
Thread 31: x[0*K + k_off + 31], x[2*K + k_off + 31], x[4*K + k_off + 31], ...
Ah! This is actually quite good. At any given iteration of the loop (fixed idx offset), threads 0-31 access x[k_off+0] through x[k_off+31] for m=0, then in the next loop iteration they access x[2*K + k_off+0] through x[2*K + k_off+31], etc.
Wait, no. In the first iteration of idx, threads access tid=0..31, which are all in m=0 (since BK=64, tid < 64 means m=0). So the first iteration of the loop for all threads accesses m=0 with k=tid. These are perfectly contiguous! Great.
In the second iteration, idx = tid + 128. For tid=0..31: idx=128..159. m = idx / 64 = 2 for idx=128..191. k = idx % 64. So for idx=128..159: m=2, k=0..31. These are contiguous in global memory? x[2*K + k_off + 0] through x[2*K + k_off + 31]. Yes! Perfectly contiguous.
So in the first loop iteration, all 128 threads access m=0 with k=0..127 (but BK=64, so this splits across m=0 with k=0..63 and m=1 with k=0..63). Wait, with BN=128 threads: tid=0..63 access m=0 with k=tid. tid=64..127 access m=1 with k=tid-64.
So the warp of threads 0-32: first 32 access m=0, k=0..31. Next warp (32-63): m=0, k=32..63. Both are contiguous in x (m=0 row).
Then threads 64..95 access m=1, k=0..31. Thread 64 accesses x[1*K + k_off + 0]. Thread 95 accesses x[1*K + k_off + 31]. This is contiguous.
But warp 2 (threads 64-95) and warp 3 (96-127) access m=1. Different warps accessing the same row is fine. Actually, all 4 warps access only 2 rows in the first iteration!
This is actually quite coalesced. The main issue is that for a single warp, not all threads access the same memory segment. But in the first iteration, warp 0 accesses the first 32 elements of row 0. That's perfectly coalesced into 1 or 2 128-byte cache lines. Warp 1 accesses the next 32 elements of row 0. Also perfectly coalesced. Warp 2 accesses the first 32 elements of row 1. Coalesced but different segment. Warp 3 accesses the next 32 elements of row 1. Coalesced.
So the first loop iteration accesses 4 segments: row 0 first half, row 0 second half, row 1 first half, row 1 second half. That's 4 memory transactions per warp? Actually, each warp accesses only one segment. Warp 0 accesses 128 bytes (32 floats). That's 2 128-byte transactions on H100 (or 1 if 256-byte L1 cache line). Perfectly coalesced.
The second loop iteration (idx += 128): tid=0 accesses idx=128 -> m=2, k=0. tid=127 accesses idx=255 -> m=3, k=63. So we access rows 2 and 3 in the same pattern.
So the loop runs M*BK / BN = 8192 / 128 = 64 iterations. Each iteration loads 2 rows (128 elements). The access pattern is perfectly coalesced within each warp.
So the simple loop is actually quite good! Let me verify:
- Each thread does 64 loads from x (8192 / 128 = 64)
- Consecutive threads access consecutive memory in each loop iteration
- All loads from x in a given iteration hit at most 4 rows, spaced 128KB apart (since K=32768, each row is 128KB)
Actually, are the loads from different rows coalesced? On modern GPUs, each warp issues a memory request for the addresses accessed by its threads. If threads access addresses within the same 128-byte segment, it's one transaction. If they span multiple segments, it's multiple transactions.
For warp 0 in iteration 0: addresses x[0*K + k_off + 0] to x[0*K + k_off + 31]. This is 128 bytes, contiguous. One transaction.
Warp 2 in iteration 0: addresses x[1*K + k_off + 0] to x[1*K + k_off + 31]. Also contiguous, one transaction. But this is a different row, 128KB away. The memory controller handles this.
So the simple loop is actually perfectly fine! And much simpler than warp-level programming. Let me stick with the simple loop.
Similarly for w_tile:
```cpp
for (int idx = tid; idx < BN * BK; idx += BN) {
int bn = idx / BK;
int k = idx % BK;
w_tile[idx] = weight[(n_start + bn) * K + k_off + k];
}
```
Consecutive threads access consecutive k values within the same bn row. This is also perfectly coalesced.
Great, the simple loops are fine.
Now let me optimize the compute loop:
```cpp
#pragma unroll
for (int k = 0; k < BK; k++) {
float w_val = w_tile[tid * BK + k];
#pragma unroll
for (int m = 0; m < M; m++) {
accum[m] += x_tile[m * BK + k] * w_val;
}
}
```
Here, x_tile[m * BK + k] for varying m at fixed k are strided by BK in shared memory. Since BK=64, they are 64 floats = 256 bytes apart. Shared memory banks are typically 4 bytes each, so 256 bytes apart means bank 0, bank 64, bank 128, etc. With 32 banks, bank 0, bank 0 (64%32=0), bank 0 (128%32=0). ALL accesses go to bank 0! Major bank conflict!
This is bad. We need to transpose the storage of x_tile or reorder the loop.
If we store x_tile as column-major (k * M + m) instead of row-major (m * BK + k):
```cpp
x_tile[k * M + m] = x[m * K + k_off + k];
```
Then in compute:
```cpp
#pragma unroll
for (int k = 0; k < BK; k++) {
float w_val = w_tile[tid * BK + k];
#pragma unroll
for (int m = 0; m < M; m++) {
accum[m] += x_tile[k * M + m] * w_val;
}
}
```
Now at fixed k, x_tile[k * M + m] for m=0,1,2,... are contiguous in shared memory. Threads in the same warp (different tid, same m iterations) access the same addresses because all threads use the same x_tile for a given (k,m). Actually, all threads access x_tile[k * M + m] simultaneously. This is a broadcast - all threads read the same address. Shared memory broadcasts are fast (handled by the broadcast mechanism).
Wait, no bank conflicts for broadcasts! All threads reading the same address is actually optimal for shared memory. But here, within a single thread, it reads x_tile[k*M + 0], x_tile[k*M + 1], ... sequentially. There's no inter-thread shared memory access pattern issue because all threads use the same x but in independent loop iterations over m.
Actually, all threads access x_tile[k * M + m] for m from 0 to M-1. For a fixed k, thread 0 reads x_tile[k*M+0], x_tile[k*M+1], ... (sequentially in loop). Thread 1 also reads x_tile[k*M+0], x_tile[k*M+1], ... but at slightly different times if the loop is unrolled. But since this is within each thread's own execution, and the memory is in shared memory, each thread reads sequentially. At any given clock cycle, all threads in a warp might be reading different m values depending on how the compiler schedules. But actually, since all threads execute the same loop, they likely access the same m values simultaneously (warp-level lockstep).
If all 32 threads in a warp access x_tile[k*M + m] for the same m but different k? No, they all use the same k in the outer loop. So for a given k, all threads access x_tile[k*M + m] where m is varying. Wait, in the inner loop over m, at iteration m=0, ALL threads access x_tile[k*M + 0]. That's a broadcast! Then at m=1, ALL threads access x_tile[k*M + 1]. Another broadcast. And so on.
With x_tile stored as k*M + m, all threads accessing the same m simultaneously means they all access the same address (since k and m are the same for all threads). That's a perfect broadcast - no bank conflicts!
But w_tile[tid * BK + k]: different threads access different rows of w_tile. At fixed k, thread 0 accesses w_tile[0*BK+k], thread 1 accesses w_tile[1*BK+k], etc. These are BK=64 floats apart. If w_tile is stored as bn*BK + k: thread 0 accesses w_tile[k], thread 1 accesses w_tile[64+k], thread 2 accesses w_tile[128+k], etc. That's strided by 64 floats = 256 bytes. Bank conflicts again!
To avoid bank conflicts for w_tile, we should store it as k * BN + bn (column-major):
```cpp
w_tile[k * BN + bn] = weight[(n_start + bn) * K + k_off + k];
```
Then access: `w_tile[k * BN + tid]`
At fixed k, threads 0..127 access w_tile[k*BN + 0], w_tile[k*BN + 1], etc. These are contiguous in shared memory! 128 threads accessing 128 consecutive floats = 512 bytes. With 32 banks, some minor conflicts might occur but much better than before. Actually, if each bank is 4 bytes, 128 consecutive floats span banks cyclically. Adjacent threads access adjacent banks, no conflicts.
But w_tile[k*BN + tid] for k varying: consecutive k values are BN floats apart. So w_tile[0*BN+tid] and w_tile[1*BN+tid] are BN floats apart. With BN=128, that's 512 bytes apart. Bank = (512/4) % 32 = 128 % 32 = 0. So consecutive k accesses hit the same bank! Bad for the outer loop if BK is large.
Hmm, we can transpose differently. But the main compute loop accesses w_val for a fixed (tid, k). Since the loop is over k, consecutive w_vals are w_tile[k*BN + tid] and w_tile[(k+1)*BN + tid]. If stored as k*BN + bn, these are BN apart in memory. So there's a sequential bank conflict in the outer loop. But shared memory latency is usually hidden by instruction-level parallelism.
Actually, since we're already likely memory-bound on loading from global memory, and shared memory has much higher bandwidth, the shared memory access pattern might not be the bottleneck.
Wait, but on H100 shared memory bandwidth is ~15 TB/s or more per SM, and we're doing a lot of compute. Let me not over-optimize unless necessary.
Actually, the key concern is that w_tile stores BN*BK = 8192 floats = 32KB. With bank conflicts, the effective bandwidth might be lower. But we have 512 iterations over K, and each iteration loads BK=64 w_vals per thread. Total shared memory reads for w: 128 threads * 512 * 64 * 4 bytes = 16MB per block. And for x_tile: 128 threads * 512 * 64 * 4 bytes = 16MB. Total shared memory bandwidth: 32MB per block. With 256 blocks, that's 8GB total shared memory traffic. At 15 TB/s per SM and 128 SMs, that's easily achievable.
Actually wait, shared memory bandwidth per SM is much higher than HBM. So shared memory is not the bottleneck.
_GLOBAL memory_ is the bottleneck. Let me recalculate global memory reads:
- x is read once: M*K = 128*32768 = 4.2M floats = 16.8MB
- weight is read once: N*K = 32768*32768 = 1.07B floats = 4.3GB
- bias is read once: N = 32768 floats = 128KB
- output is written once: M = 128 floats = 512 bytes
Total global memory reads: ~4.3GB + 16.8MB + 128KB ≈ 4.32GB.
Time at 3.35 TB/s: ~1.29ms.
Total FLOPs: 2 * M*N*K + M*N (for sigmoid, but that's negligible) ≈ 275B FLOPs.
Time at 67 TFLOPS: ~4.1ms.
So this is compute-bound IF global memory can be read fast enough and cached. But for the tiled approach, each weight tile is reused across all M batch elements within the block. But different blocks read different tiles of weight. Since there are 256 blocks and weight is 4GB, and L2 cache is ~50MB, at most about 1% of weight can be in L2. So almost all weight reads go to HBM.
Weight is read exactly once by each block. Wait, no! In my design, each block handles BN=128 different n values. A weight element weight[n, k] is needed only by the block that handles that n value. Since each weight row is assigned to exactly one block, each weight element is read exactly once total across all blocks.
Wait, is that right? Yes! Weight row n is read by exactly one block (the block with n_start ≤ n < n_start + BN). So total weight reads across all blocks = N * K = 4.3GB. That's optimal - we read weight exactly once.
Similarly, x is read by all blocks. Each block reads all of x. But x is only 16.8MB! So total x reads = 256 blocks * 16.8MB = 4.3GB. Wait, that's also 4.3GB! Because each of 256 blocks reads all of x (128 rows * 32768 cols = 4.2M floats = 16.8MB). 256 * 16.8MB = 4.3GB.
So total memory reads: 4.3GB (weight) + 4.3GB (x) = 8.6GB! At 3.35 TB/s, that's 2.6ms.
Add to that compute of 275B FLOPs at 67 TFLOPS = 4.1ms.
If we can perfectly overlap memory and compute, we'd be compute-bound at 4.1ms. But realistically some memory traffic will stall compute. Let me see if cuBLAS does better.
In cuBLAS GEMM for C = A @ B:
- A (x) is M×K, read once: 16.8MB
- B (weight.t()) is K×N, read once: 4.3GB
- C is M×N, written once: 16.8MB
Total memory: ~4.3GB reads + 16.8MB write = 4.32GB.
Time at 3.35 TB/s: ~1.3ms.
Wait, but cuBLAS reads weight.t() which is K×N. In row-major, weight.t()[k, n] = weight[n, k]. So reading weight.t() by tiles means reading weight with a specific access pattern. But weight[n, k] stored as (N, K) row-major. For a tile of B that spans k0..k1 and n0..n1, the elements are weight[n, k] for n in [n0,n1) and k in [k0,k1). In global memory, for fixed n, varying k is contiguous. For fixed k, varying n is strided by K floats.
cuBLAS handles this by clever tiling. But the total memory read is still ~4.3GB because each weight element is read once.
In our kernel, we also read weight once. And x is read 256 times (once per block) because x is small enough that it might stay in L2? 16.8MB * 256 = 4.3GB total. Wait, but x could be cached in L2! L2 is 50MB. x is only 16.8MB. So if x stays in L2 after the first few blocks, subsequent reads are from L2, not HBM. L2 bandwidth is much higher. So actual HBM traffic for x might be just 16.8MB.
Hmm, but blocks run in parallel. Multiple blocks might read x simultaneously. Still, once x is in L2, subsequent reads of x by later blocks will be L2 hits. So effective HBM reads might be close to 4.3GB (weight) + 16.8MB (x, once) + 128KB (bias, once) ≈ 4.32GB.
That's actually very similar to cuBLAS! So our kernel and cuBLAS would have similar memory traffic.
But cuBLAS is much more optimized for compute - better instruction scheduling, register usage, etc. So can our hand-written kernel beat cuBLAS? Unlikely, unless cuBLAS has some inefficiency for this small-M case.
Wait, but our kernel fuses sigmoid and sum! That saves writing and reading the intermediate (M, N) matrix. The reference writes 16.8MB for the linear output, then another 16.8MB for sigmoid (if in-place, actually sigmoid reads and writes in-place, so just 16.8MB), then sum reads 16.8MB.
Wait, in PyTorch:
1. linear writes C (M, N) = 16.8MB
2. sigmoid reads C and writes C (in-place) = 16.8MB read + 16.8MB write = 33.6MB
3. sum reads C and writes output (M, 1) = 16.8MB read + 0.5KB write
Total additional memory for epilogues: 33.6MB + 16.8MB = 50.4MB.
So our kernel saves 50.4MB of memory traffic compared to reference. But the matmul itself does 4.3GB. So the savings are about 1.2% of total memory. That translates to maybe 0.05ms saved on a 4ms kernel. Barely measurable.
But wait - there's something else! In our kernel, after computing the dot products, we don't write anything to global memory except the final 128 floats. The reference writes and reads 16.8MB multiple times. For a problem that's potentially memory-bound (if cuBLAS is limited by HBM bandwidth), avoiding these intermediate writes could help.
Actually, is the reference compute-bound or memory-bound?
- FLOPs: 275B
- Memory: 4.3GB (matmul) + 50MB (epilogues and intermediate)
- Arithmetic intensity: 275B / 4.35GB = 63 FLOPs/byte
- H100: 67 TFLOPS / 3.35 TB/s = 20 FLOPs/byte
Since 63 > 20, this is compute-bound for the matmul. The epilogues are negligible.
So fusing won't help much for a compute-bound matmul.
Hmm, unless PyTorch's matmul is not hitting peak FLOPs. For M=128, N=32768, K=32768, this is a somewhat small M. cuBLAS has overhead and might not reach full utilization for M=128.
Actually, on H100 with 128 SMs, M=128 means each SM handles at most 1 row. But N=32768 and K=32768 are large, so each row has a lot of work. I think cuBLAS should handle this well, but maybe not perfectly.
What about using Tensor Cores? If we can write a kernel that uses Tensor Cores for the matmul and gets significantly faster, even if accuracy is slightly degraded, might we still pass? The tolerance is 1e-4. Tensor cores in FP32 accumulation mode (e.g., FP16 input × FP16 input with FP32 accumulation) might have accuracy issues. But TF32 with FP32 accumulation... TF32 has 10 mantissa bits. For a dot product of 32768 values, the accumulated error might be larger than 1e-4.
Actually let me calculate: TF32 has about 3 decimal digits of precision. For dot products of length 32768, the error could accumulate significantly. But if the values are random [0,1], the dot product is around 16384. The absolute error from TF32 might be ~16384 * 1e-3 = 16. That's way larger than 1e-4 tolerance. So TF32 is definitely out.
What about FP16 with FP32 accumulation? Similar issue. The inputs are [0,1], so FP16 can represent them well. But the products x*w are ~1. Sum of 32768 such products ~16384. The rounding error per FMA is ~0.00003 in FP32 accumulation... actually for FP16×FP16 with FP32 accumulation, each product is rounded to FP16 (since inputs are FP16), so error per product ~0.00003 * 0.00003? No, FP16 has ~3e-4 relative precision. Product of two [0,1] values in FP16 has error ~3e-4 * 1 = 3e-4 absolute. Summing 32768 such products: if errors are random, variance grows, but expected absolute error might be 32768 * 3e-4 = ~10. That's too much.
So we need FP32 accumulation. That limits us to 67 TFLOPS.
Wait, wait. The H100 spec says FP32 TFLOPS: 67. But TF32 Tensor Core TFLOPS: 989. TF32 _accumulates_ in FP32. Each FMA is: multiply two TF32 numbers, round to TF32(? No, the hardware does the multiply at higher precision, then accumulates to FP32). Actually TF32 tensor cores do: A and B are TF32 format (1 sign, 8 exponent, 10 mantissa bits). The multiplication produces an FP32 result (or close to it), and accumulation is in FP32.
But the issue is that inputs are rounded to TF32 precision (10 mantissa bits) at the start. So even though accumulation is FP32, the per-element precision is limited to ~1e-3 relative. For dot products, the error accumulates.
Actually, let me reconsider whether TF32 could pass. PyTorch might already be using TF32 for this matmul. If so, the reference output already has TF32-level accuracy, and we just need to match it. If the reference doesn't use TF32, then using TF32 might violate the numerical tolerance.
I should first try to just write the best possible FP32 kernel and see if it wins. Maybe the reference implementation has some inefficiency that I'm not accounting for.
Let me also consider: what if M=128 is small enough that the kernel launch overhead is significant? The reference does 3 kernel launches (linear, sigmoid, sum). Our fused kernel does 1 launch. For a 0.5ms kernel, saving 2 launches might be 10-20 microseconds. Not much.
But wait - in PyTorch, `nn.Linear` might not call `addmm` with optimal settings. Let me verify what actually happens.
In PyTorch `nn.Linear`:
```python
def forward(self, input):
return F.linear(input, self.weight, self.bias)
```
`F.linear`:
```python
if input.dim() == 2 and bias is not None:
# fused op is marginally faster
ret = torch.addmm(bias, input, weight.t())
```
So yes, for 2D input, it uses `torch.addmm(bias, input, weight.t())` which is a single cuBLAS call that computes `beta * bias + alpha * (input @ weight.t())`. This is highly optimized.
So the reference is: one cuBLAS call + one elementwise sigmoid + one reduction sum.
OK, given all this, maybe I should write the kernel and see. Let me also consider another optimization: since M=128, N=32768, K=32768, and we know this at compile time, we might be able to use the `mma.sync` instructions or the `wmma` API to get tensor core acceleration without going through cuBLAS. But that would require writing a fairly complex kernel.
Actually, H100 has a lot of SMs. With M=128 and N=32768, even if each SM handles 1 row and processes N elements with high parallelism, there are only 128 rows for 128 SMs. That's perfect SM utilization!
But the issue is that each row has K=32768, and to compute one output element y[m,n], we need to do a dot product. The standard GEMM maps well to this.
Another idea: use a persistent kernel style where threads stay resident and process multiple n values. But that's what cuBLAS already does.
OK, I'll write the kernel as described and test it. If it's not faster, I'll try other optimizations.
Also, I can try using `float4` loads for x_tile and w_tile to improve global memory bandwidth.
Let me think about float4:
- x_tile has M*BK = 8192 floats. Float4 loads mean 2048 float4 loads.
- With 128 threads, each thread does 16 float4 loads.
- w_tile similarly.
For x: x[m * K + k_off + k] with k increasing by 4. This is contiguous. Great.
For w: weight[n * K + k_off + k] with k increasing by 4. Also contiguous. Great.
Using float4 will significantly improve global memory throughput.
Let me write the kernel with float4. But I need to be careful about alignment. K=32768 is divisible by 4, so k_off is always divisible by 4 (since BK should be a multiple of 4). I'll use BK=64, which is divisible by 4.
For x_tile and w_tile in shared memory, if I use float4 loads, I need to access shared memory as float4 too. Or I can load into a local Float4 and then write to shared memory as individual floats. Loading to local first is easier.
Actually, I can declare x_tile as float4 array: `__shared__ float4 x_tile4[M * BK / 4];` but then accessing individual floats for computation requires type punning. Let me just load with float4 and store individual floats:
```cpp
float4 x4 = *(float4*)(x_ptr + offset);
x_tile[offset] = x4.x;
x_tile[offset+1] = x4.y;
x_tile[offset+2] = x4.z;
x_tile[offset+3] = x4.w;
```
But this is messy. Alternatively, just load floats without vectorization and let the compiler handle it. The memory controller on H100 is smart enough to coalesce 128 concurrent loads into a few transactions. The actual bandwidth might already be near peak.
OK let me just write the first version without float4 and see how it performs. I can optimize later.
One more thought: can we compute the sum without atomics by having blocks write partial sums and then summing with a second kernel? Let's do the atomic version first for simplicity.
Wait, actually for the output, each block needs to add its partial sum. With 256 blocks, we can allocate a buffer of M * 256 floats, have each block write there, then reduce with a simple kernel. This avoids atomic contention. The buffer size is 128K bytes which is tiny.
Let's do this: the main kernel writes to `partial_sums[blockIdx.x * M + m]`. Then a tiny reduce kernel sums across blocks. Or we can even use PyTorch to do the reduction after the custom kernel returns partial_sums.
In the Python wrapper:
```python
partial_sums = torch.zeros(N//BN, M, device=x.device)
custom_kernel(x, weight, bias, partial_sums)
output = partial_sums.sum(dim=0, keepdim=True).t()
```
But using PyTorch's `sum` adds overhead. A custom tiny reduction kernel is better:
```cuda
__global__ void reduce_kernel(const float* partial_sums, float* output) {
// One block, M threads
int m = threadIdx.x;
float sum = 0;
for (int b = 0; b < num_blocks; b++) {
sum += partial_sums[b * M + m];
}
output[m] = sum;
}
```
This uses 128 threads doing 256 loads each = 32768 loads. Trivial.
But launching a second kernel has overhead. Let's see if atomics are acceptable. For 128 atomic adds per block to distinct addresses, that's 32768 atomic operations. Actually, each m address gets 256 atomic adds. On H100, `atomicAdd` to different addresses within a warp is coalesced where possible. But here, all threads in a block might do atomics at slightly different times. And across 256 blocks, the atomics to the same 128 addresses might have some contention. But honestly, 256 atomic adds per address over the lifetime of the kernel is very manageable. With 128 addresses and 128 SMs, at peak there might be ~1-2 blocks per SM running concurrently, so each address gets at most ~2 concurrent atomic adds. That's nothing for H100.
I'll use atomics for simplicity, but I need to zero-initialize the output first.
Actually, another option: since each block only writes atomics to M=128 distinct addresses, and each block handles BN=128 different n values, we could have a shared memory array `block_output[M]`, initialize to 0, accumulate the warp sums, and then have `threadIdx.x < M` write `atomicAdd(&output[m], block_output[m])`. Since M=128 = block size, exactly one thread per output.
Wait, I'd need all warps to write to `block_output`. Let me design the reduction properly.
OK, let me write the code. I'll start with a simpler reduction approach:
```cpp
// After compute loop, each thread has M accumulators
// Compute sigmoid and thread-local sum
float thread_local_sum[M];
for (int m = 0; m < M; m++) {
float dot = accum[m] + bias[n];
thread_local_sum[m] = 1.0f / (1.0f + expf(-dot));
}
__shared__ float block_sums[M];
if (threadIdx.x == 0) {
for (int m = 0; m < M; m++) block_sums[m] = 0.0f;
}
__syncthreads();
for (int m = 0; m < M; m++) {
atomicAdd(&block_sums[m], thread_local_sum[m]);
}
__syncthreads();
if (threadIdx.x == 0) {
for (int m = 0; m < M; m++) {
atomicAdd(&output[m], block_sums[m]);
}
}
```
Wait, `atomicAdd` to shared memory is fine and much faster than global memory. We can use `atomicAdd(&block_sums[m], thread_local_sum[m])` for all threads. But there might be shared memory bank conflicts if multiple threads hit the same bank. Here, all threads writing to block_sums[m] for the same m but different m across threads would mean thread 0 writes to block_sums[0], thread 1 to block_sums[1], etc. These are consecutive in shared memory, so different banks (assuming 32 banks). So no conflicts! Perfect.
Actually wait - each thread iterates m from 0 to 127. So at step m=0, all 128 threads write to block_sums[0] atomically. That's a huge shared memory conflict! All threads hitting the same address. For 128 threads, `atomicAdd` to the same address would serialize. And we'd do this 128 times (for m=0 to m=127). The total serialization would be 128 * 128 = 16384 sequential operations. That's bad.
Better approach: have each thread accumulate its own partial sum in local registers, then do warp-level reduction, then atomicAdd block-level results.
```cpp
// Each thread computes thread_sum[m] = sigmoid(dot(x[m], W[n]) + bias[n])
float thread_sums[M];
for (int m = 0; m < M; m++) {
thread_sums[m] = 1.0f / (1.0f + expf(-(accum[m] + bias[n])));
}
// Warp-level reduction
for (int m = 0; m < M; m++) {
float val = thread_sums[m];
for (int offset = warpSize/2; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (threadIdx.x % warpSize == 0) {
warp_sums[(threadIdx.x / warpSize) * M + m] = val;
}
}
__syncthreads();
// Inter-warp reduction (only warp 0)
if (threadIdx.x < warpSize) {
for (int m = 0; m < M; m++) {
float val = 0.0f;
int num_warps = blockDim.x / warpSize; // 4
for (int w = 0; w < num_warps; w++) {
val += warp_sums[w * M + m];
}
if (threadIdx.x == 0) {
atomicAdd(&output[m], val);
}
}
}
```
Wait, `for (int w = 0; w < num_warps; w++)` only by threads in warp 0. That's fine.
warp_sums needs to be in shared memory: `__shared__ float warp_sums[4 * M];` = 512 floats = 2KB. Fine.
This is much better. Only 4 atomic adds per m (one per warp, but actually only 1 atomic after inter-warp reduction). Total atomics: 128.
Now, one issue: bias[n] is accessed once per thread. bias[n] where n = n_start + threadIdx.x. So threads in a warp access consecutive bias elements. That's coalesced.
Another performance consideration: the inner compute loop over m with M=128. This is a fairly large loop. With accum[128] in registers, that's 128 * 4 = 512 bytes of registers per thread. The compiler should unroll this.
But wait, a thread with 128 float accumulators might run out of registers. Let me calculate: 128 accumulators + a few temporaries = ~132 registers. H100 max is 255, so this is fine. But occupancy might be limited if all threads use many registers.
With 64K registers per SM and 128 threads per block, using 132 registers per thread requires ~16.8K registers per block. An SM can handle up to 64K / 16.8K = ~3.8 blocks concurrently. With 64KB shared memory per block and 228KB total, that's ~3.5 blocks per SM. So we might get 3 blocks per SM. With 128 SMs, that's 384 blocks concurrently. But we only have 256 blocks. So each SM gets about 2 blocks, and we have 128 SMs * 2 = 256. That works out perfectly!
Actually, 256 blocks / 128 SMs = 2 blocks per SM. With 64KB shared memory per block, 2 blocks use 128KB, well under 228KB. With ~132 registers per thread, 2 blocks * 128 threads * 132 registers = 33.8K registers, under 64K.
So 2 blocks per SM = 256 threads per SM. H100 has 128 cores per SM? H100 has 128 FP32 CUDA cores per SM? Let me think... Actually H100 SM has 128 FP32 cores. With 256 threads (2 blocks * 128 threads, but actually threads run in warps), 256 threads = 8 warps. Each SM can handle up to 64 warps. So occupancy is 8/64 = 12.5%. That's quite low!
Low occupancy means fewer warps to hide latency. But our kernel is compute-bound with long-running loops, so latency hiding might not be critical. Still, we might want to increase occupancy.
To increase occupancy, we need either:
1. More threads per block (but then shared memory per block increases if we keep same tiles)
2. Larger BN to amortize overhead
3. More blocks per SM by reducing shared memory/register usage
One option: increase BN to 256. Then w_tile = 256 * 64 * 4 = 64KB. x_tile = 32KB. Total = 96KB. Still under 128KB per block. With 228KB per SM, we can fit 2 blocks (192KB).
But then blockDim.x = 256 threads. Accumulators per thread = 128 floats. Registers = ~132 * 256 = 33.7K per block. 2 blocks = 67.5K > 64K limit! So we can only run 1 block per SM if we have 256 threads. But 1 block with 256 threads also gives 8 warps per SM (assuming 32 threads per warp). Wait, 256 threads = 8 warps. Still low occupancy.
Hmm, what if we process multiple n values per thread? Say BN=256 but blockDim.x=128, so each thread handles 2 n values. Then accumulators per thread = 2 * 128 = 256 floats = 1024 bytes of registers. ~260 registers. That's near the limit and might reduce occupancy further.
What if we reverse the loop? Instead of each thread handling one n and all m, each thread handles one m and several n values?
Actually, M=128 is small. We could process m across threads and n across blocks. But then x broadcast pattern changes.
Alternative: process with 2D thread blocks. blockDim.x for n, blockDim.y for m. Say 32x4 = 128 threads. Each thread computes one (m, n) output. That requires just 1 accumulator per thread! Registers drop significantly.
But then for reduction, we'd need to sum across all threads for each m. With 128 n values per block and 4 m values per thread... hmm.
Wait, actually if we use blockDim=(BN, M/BM) where BM is batch elements per block, but M=128 and if BM=1, then blockDim.y=128, giving 128*128=16384 threads. That's way too many.
Actually, we can split M across threads. But M=128 and BN=128 gives 16384 outputs per block. With 128 threads, each thread handles 128 outputs. So each thread would need 128 accumulators. Same issue.
What if we process M in a coarser way? Since M=128 is the total batch size, we can't reduce it.
Wait, what if we use fewer registers by not fully unrolling? Or use float2/float4 for accumulators? That doesn't reduce register count.
Another idea: instead of computing all M=128 outputs per thread, have the block compute all M outputs but use shared memory for partial accumulations. For example, a thread computes partial dot products and accumulates into shared memory. But then we'd have shared memory conflicts.
Actually, the register approach with 128 floats per thread is probably fine. 128 registers per thread means with 128 threads per block, 128*128 = 16K registers. With 64K per SM, we can have 4 blocks per SM! But we only launch 256 blocks and have 128 SMs. So 2 blocks per SM. Still only 4 warps per SM if each block has 4 warps (128 threads).
Wait, 128 threads = 4 warps of 32 threads. 2 blocks per SM = 8 warps per SM. H100 SM can hold 64 warps. 8/64 = 12.5% occupancy. That's the same as before. But if we can fit 4 blocks per SM, that's 16 warps = 25% occupancy. Still not great for latency hiding, but for a compute-bound kernel, it might not matter.
Can we fit 4 blocks per SM?
- Shared memory: 4 * 64KB = 256KB > 228KB. No.
- With TN=64 instead of 128: shared mem = 32KB + 32KB = 64KB. 4 blocks = 256KB. Still too much.
- With TN=64 and BK=32: x_tile = 16KB, w_tile = 8KB. Total = 24KB. 4 blocks = 96KB. Under 228KB!
- But then K tiles = 32768/32 = 1024. More iterations.
With TN=64, BK=32, block=64 threads:
- x_tile = 128*32*4 = 16384 bytes
- w_tile = 64*32*4 = 8192 bytes
- Total = 24KB per block
- 4 blocks per SM = 96KB
- Registers per block: 64 threads * ~132 regs = 8448. 4 blocks = 33.8K. Under 64K!
- Threads per SM: 4 * 64 = 256. 8 warps. Still only 12.5% of max warps (64).
Hmm, the fundamental issue is that M=128 requires some minimum number of registers or shared memory. But 8 warps per SM for a compute-bound kernel might be OK because the ALUs can stay busy.
Actually, let me look at this from a different angle. If we have 128 SMs and 256 blocks, 2 blocks per SM. Each block computes for 128 n-values across all 128 m-values. FMAs per block: 128*128*32768 = ~536M. At 67 TFLOPS across 128 SMs, each SM gets ~523M FLOPs. At 67 TFLOPS / 128 SMs = ~523 GFLOPS per SM. Time per SM: 536M / 523G = ~1.02ms. With 2 blocks per SM, each block takes ~0.5ms? No wait, the SMs run blocks concurrently. Total time for the kernel would be around the time per block if all SMs are fully utilized with 2 blocks each.
Actually, total time ≈ total_FLOPs / (SMs * FLOPS_per_SM) = 275B / (128 * 523G) = 275B / 67T = 4.1ms. So the occupancy issue doesn't hurt the compute throughput if the warps are fully busy.
OK, I think my initial design is solid. Let me write it and test.
One more optimization: I can increase BN to 256 to reduce the number of blocks to 128. Then we have exactly 1 block per SM. With BN=256 and BK=64:
- w_tile = 256*64*4 = 64KB
- x_tile = 128*64*4 = 32KB
- Total = 96KB per block. Under 128KB per block, and under 228KB per SM.
- Threads = 256 per block. 256 threads = 8 warps.
- Registers: 256 threads * 132 regs = 33.7K. Under 64K.
- Blocks: 128 (exactly one per SM!)
This is actually very clean! 1 block per SM, 8 warps per SM. H100 should run this well. And we don't need cross-block synchronization or atomics from multiple blocks to the same output (but we still have 128 blocks with only 128 SMs, so atomics are still needed).
Wait, with 256 threads per block, and each thread handling one n value and all M=128 m values, we have 256 n values per block. N=32768, so 32768/256 = 128 blocks. Perfect!
Let me redo the kernel with BN=256, threads=256:
```cuda
#define M 128
#define K 32768
#define N 32768
#define BK 64
#define BN 256
#define NUM_K_TILES (K / BK) // 512
#define NUM_BLOCKS (N / BN) // 128
__global__ void fused_kernel(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output
) {
__shared__ float x_tile[M * BK];
__shared__ float w_tile[BN * BK];
__shared__ float warp_sums[(BN / 32) * M]; // 8 * 128 = 1024 floats = 4KB
int n_start = blockIdx.x * BN;
int tid = threadIdx.x;
int n = n_start + tid;
int lane_id = tid % 32;
int warp_id = tid / 32;
int num_warps = BN / 32; // 8
float accum[M];
#pragma unroll
for (int m = 0; m < M; m++) accum[m] = 0.0f;
for (int kt = 0; kt < NUM_K_TILES; kt++) {
int k_off = kt * BK;
// Load x_tile: M * BK = 8192 floats
for (int idx = tid; idx < M * BK; idx += BN) {
int m = idx / BK;
int k = idx % BK;
x_tile[idx] = x[m * K + k_off + k];
}
// Load w_tile: BN * BK = 16384 floats
for (int idx = tid; idx < BN * BK; idx += BN) {
int bn = idx / BK;
int k = idx % BK;
w_tile[idx] = weight[(n_start + bn) * K + k_off + k];
}
__syncthreads();
// Compute partial dots
#pragma unroll
for (int k = 0; k < BK; k++) {
float w_val = w_tile[tid * BK + k];
#pragma unroll
for (int m = 0; m < M; m++) {
accum[m] += x_tile[m * BK + k] * w_val;
}
}
__syncthreads();
}
// Apply sigmoid and prepare for reduction
float thread_sums[M];
#pragma unroll
for (int m = 0; m < M; m++) {
float val = accum[m] + bias[n];
thread_sums[m] = 1.0f / (1.0f + expf(-val));
}
// Warp-level reduction for each m
#pragma unroll
for (int m = 0; m < M; m++) {
float val = thread_sums[m];
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (lane_id == 0) {
warp_sums[warp_id * M + m] = val;
}
}
__syncthreads();
// Inter-warp reduction and atomic add
if (tid < 32) {
#pragma unroll
for (int m = 0; m < M; m++) {
float val = 0.0f;
#pragma unroll
for (int w = 0; w < num_warps; w++) {
val += warp_sums[w * M + m];
}
if (tid == 0) {
atomicAdd(&output[m], val);
}
}
}
}
```
Wait, `__shfl_down_sync` offsets: for warp size 32, we need offsets 16, 8, 4, 2, 1. Let me write it as:
```cpp
for (int offset = 16; offset > 0; offset >>= 1) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
```
This reduces 32 values to 1.
For the inter-warp reduction: we have 8 warp sums. I used `tid < 32` (warp 0) to read all 8 warp sums and add them. One thread (tid==0) does the atomicAdd. This is fine.
But wait - the `for (int w = 0; w < num_warps; w++)` loop reads from shared memory with the same address for all threads in warp 0. Bank conflicts? warp_sums[w * M + m] for fixed m but different w. Adjacent threads in warp 0 (tid=0..31) all read the same w=0 address in the first iteration. That's a broadcast, fine. But they do this sequentially for w=0, w=1, etc. Actually all threads read the same elements, so it's fine. The compiler could unroll this.
Wait, all 32 threads in warp 0 are redundant. Only one thread needs to do the sum and atomic. But since warp 0 is idle anyway, it's fine. The shared memory reads by all threads are broadcasts.
Actually, I can simplify: have `tid == 0` do all the summing and atomicAdd:
```cpp
if (tid == 0) {
for (int m = 0; m < M; m++) {
float val = 0.0f;
for (int w = 0; w < num_warps; w++) {
val += warp_sums[w * M + m];
}
atomicAdd(&output[m], val);
}
}
```
This is even simpler. But having just 1 thread reading shared memory might be slightly worse than having all 32 threads cooperatively load. But with only 8 reads per m, and 128 m values, that's 1024 reads total. A single thread can do this in microseconds. It's fine.
OK so tid == 0 does the final reduction and atomic add.
Now for initialization: output must be zeroed before the kernel runs. I'll do this in the Python code with `torch.zeros(M, device=x.device, dtype=torch.float32)`.
Let me also make sure I handle the grid/block dimensions correctly:
```python
threads = 256
blocks = 128
kernel[blocks, threads](x, weight, bias, output)
```
Now let me think about whether I need to synchronize between blocks. Since we're using atomics, no need.
One subtle issue: `bias[n]` where `n = n_start + tid`. The bias tensor has shape (N,). `bias[n]` accesses `bias[n_start + tid]`. For threads 0..255, this accesses contiguous elements. Since data is float and contiguous, this is perfectly coalesced.
Also, weight access: `weight[(n_start + bn) * K + k_off + k]`. For the loading loop, as idx goes from tid to BN*BK step BN, bn = idx / BK goes from 0 to BN-1, and k = idx % BK goes from 0 to BK-1. For fixed bn and varying k... actually as idx increments by BN=256, bn increases by 256/BK = 4. So bn goes 0, 4, 8, ... Wait, BK=64. idx=0: bn=0, k=0. idx=256: bn=4, k=0. idx=512: bn=8, k=0. So thread tid loads every 4th bn, and within each bn, loads all k values sequentially.
Wait, let me recalculate:
idx = tid + i * BN
bn = idx / BK = (tid + i*BN) / BK
k = idx % BK = (tid + i*BN) % BK
With BN=256, BK=64:
For tid=0: idx = 0, 256, 512, 768, ...
bn = 0/64=0, 256/64=4, 512/64=8, 768/64=12, ...
k = 0, 0, 0, 0, ...
So thread 0 loads weight[n_start+0, k_off+0], weight[n_start+4, k_off+0], weight[n_start+8, k_off+0], ...
These are strided by 4 * K = 131072 floats = 512KB. Not contiguous! But across the warp: thread 0 loads bn=0, thread 1 loads bn=0, k=1? Let's check:
tid=1: idx = 1, 257, 513, 769, ...
bn = 1/64=0, 257/64=4, 513/64=8, ...
k = 1, 1, 1, 1, ...
So thread 1 loads weight[n_start+0, k_off+1], weight[n_start+4, k_off+1], etc.
Threads 0-31 in a warp: all load bn=0 for idx < 64, then bn=4 for idx=256..319, etc.
For idx=0..31: bn=0, k=idx. So threads 0-31 load weight[n_start, k_off+0] through weight[n_start, k_off+31]. These are contiguous! Perfectly coalesced.
Then for idx=256..287: bn=4, k=idx%64=0..31. Threads 0-31 load weight[n_start+4, k_off+0] through weight[n_start+4, k_off+31]. Also contiguous and coalesced.
So within each loop iteration i, all threads access contiguous memory. Between loop iterations, they jump to bn+4. This is still fine because each access is coalesced within the warp.
But wait, each thread makes (BN*BK)/BN = BK = 64 iterations of the load loop. With 256 threads, each thread loads BN*BK/BN = BK = 64 elements. Actually, the loop `for (int idx = tid; idx < BN * BK; idx += BN)` runs (BN*BK)/BN = BK = 64 iterations for each thread.
In iteration i=0: idx = tid. bn = tid/BK = 0 for tid < 64, 1 for 64 ≤ tid < 128, etc. k = tid % 64.
For warp 0 (tid=0..31): bn=0, k=0..31. Accesses weight[n_start, k_off + 0..31]. Coalesced.
For warp 1 (tid=32..63): bn=0, k=32..63. Accesses weight[n_start, k_off + 32..63]. Coalesced.
For warp 2 (tid=64..95): bn=1, k=0..31. Accesses weight[n_start+1, k_off + 0..31]. Coalesced.
...
For warp 7 (tid=224..255): bn=3, k=32..63. Accesses weight[n_start+3, k_off + 32..63]. Coalesced.
In iteration i=1: idx = tid + 256. bn = (tid+256)/64 = 4 + tid/64. k = tid % 64.
Warp 0: bn=4, k=0..31. weight[n_start+4, k_off + 0..31]. Coalesced.
This pattern is great! Each warp always accesses a contiguous 32-float segment. Perfect coalescing.
Then in iteration i=2: idx = tid + 512. bn = 8 + tid/64. And so on up to bn=252 + tid/64.
So the w_tile loading is very efficient.
Similarly for x_tile loading:
`for (int idx = tid; idx < M * BK; idx += BN)`
M*BK = 8192. BN = 256. Each thread loads 8192/256 = 32 elements.
Iteration i=0: idx = tid. m = idx/BK = tid/64. k = tid%64.
tid=0..31: m=0, k=0..31. x[0, k_off + 0..31]. Coalesced.
tid=32..63: m=0, k=32..63. x[0, k_off + 32..63]. Coalesced.
tid=64..95: m=1, k=0..31. x[1, k_off + 0..31]. Coalesced.
...
tid=224..255: m=3, k=32..63. x[3, k_off + 32..63]. Coalesced.
Iteration i=1: idx = tid + 256. m = (tid+256)/64 = 4 + tid/64.
tid=0..31: m=4, k=0..31. x[4, k_off + 0..31]. Coalesced.
This is perfect! Each warp accesses contiguous memory.
Great, the memory loading is highly optimized.
Now for the compute loop:
```cpp
for (int k = 0; k < BK; k++) {
float w_val = w_tile[tid * BK + k];
for (int m = 0; m < M; m++) {
accum[m] += x_tile[m * BK + k] * w_val;
}
}
```
`w_tile[tid * BK + k]` for fixed k and varying tid: tid*BK means addresses 0, 64, 128, 192, 256, ... for tid=0,1,2,3,... These are 64 floats = 256 bytes apart. With 32 banks, 256 bytes / 4 bytes = 64 banks apart. 64 % 32 = 0. So all threads access bank 0! Massive bank conflicts!
Oops, this is bad. I need to fix the w_tile layout.
If I store w_tile as `w_tile[k * BN + bn]` instead of `w_tile[bn * BK + k]`:
- Load: `w_tile[k * BN + tid]` (for thread tid computing bn=tid)
- At fixed k, threads access k*BN + 0, k*BN + 1, k*BN + 2, ... which are contiguous. No bank conflicts!
But then loading from global memory becomes:
```cpp
for (int idx = tid; idx < BN * BK; idx += BN) {
int bn = idx % BN; // Wait, this doesn't work with the current iteration pattern
```
Hmm, if I store w_tile as `w_tile[k * BN + bn]`, then to load from global memory:
The natural iteration gives bn = idx / BK, k = idx % BK.
I want to write to `w_tile[k * BN + bn]`.
So:
```cpp
w_tile[(idx % BK) * BN + (idx / BK)] = weight[...];
```
This is a scatter into shared memory. But that's fine for loading. The access pattern:
For idx = 0: bn=0, k=0 -> w_tile[0*256+0] = w_tile[0]
For idx = 1: bn=0, k=1 -> w_tile[1*256+0] = w_tile[256]
For idx = 2: bn=0, k=2 -> w_tile[512]
This means thread 0 writes to w_tile[0], w_tile[256], w_tile[512], ... (strided by BN in shared memory). But threads in a warp write adjacent idx values, so:
Thread 0: w_tile[0], w_tile[256], w_tile[512], ...
Thread 1: w_tile[1], w_tile[257], w_tile[513], ...
These are adjacent in shared memory for the same k offset. For idx=0..31: these are all bn=0 (since 0..31 < BK=64), k=0..31. So w_tile[k*256+0] for k=0..31. That's w_tile[0], w_tile[256], w_tile[512], ... w_tile[7936]. These are 256 floats = 1024 bytes apart. 1024/4 = 256 banks apart. 256 % 32 = 0. All bank 0! Bank conflicts again!
To avoid bank conflicts during loading, we need the threads to write to adjacent banks.
Let me reconsider. The natural memory layout for w is row-major: weight[n, k] for k in [0, K) is contiguous. We load a tile: for bn in [0, BN), k in [0, BK), weight[n_start+bn, k_off+k] is at offset (n_start+bn)*K + k_off + k.
For adjacent threads with adjacent tid, the global memory accesses are:
tid=0: (n_start+0)*K + k_off + 0
tid=1: (n_start+0)*K + k_off + 1
...
These are contiguous.
When writing to shared memory, if we want adjacent threads to write to adjacent banks, the shared memory offset should be equal to tid (or proportional).
But our calculation: idx = tid + i*BN. bn = idx / BK, k = idx % BK.
For tid=0..31 and BK=64: bn=0, k=tid. So we want to write weight[n_start, k_off+tid] to some shared memory location.
In the compute phase, we access by (bn=tid, k). If we access `w_tile[bn * BK + k]` = `w_tile[tid * BK + k]`, then for fixed k, threads access tid*BK apart - bad.
If we access `w_tile[k * (BN+pad) + bn]`, then for fixed k, threads access k*stride + tid, which are adjacent. Good for compute.
To load: we need to write `w_tile[k * (BN+pad) + bn]`.
For the load loop: bn = idx / BK, k = idx % BK.
So `w_tile[k * (BN+pad) + bn] = weight[...]`.
For idx = tid: bn = tid / BK = 0 (if tid < 64), k = tid.
So thread tid writes to w_tile[tid] if bn=0. Wait: k*(BN+pad) + bn = tid * 257 + 0. For tid=0: w_tile[0]. tid=1: w_tile[257]. tid=2: w_tile[514].
Adjacent threads write 257 floats apart = 1028 bytes. 1028/4 = 257 banks apart. 257 % 32 = 1. So consecutive threads hit consecutive banks! Perfect!
Wait, for BN=256, let BN_PAD = 257 (or 260 to be safe, or 264 to align).
Then `w_tile[k * BN_PAD + bn]`.
For idx with bn=0, k=tid (when tid < 64):
- Thread 0: w_tile[0*257 + 0] = w_tile[0]
- Thread 1: w_tile[1*257 + 0] = w_tile[257]
- Thread 2: w_tile[2*257 + 0] = w_tile[514]
These are 257 floats apart = 1028 bytes. Bank = (1028 / 4) % 32 = 257 % 32 = 1. So adjacent threads write to banks 0, 1, 2, ... Great! No bank conflicts.
For threads tid=64..127: bn=1, k=tid%64=tid-64.
- Thread 64: w_tile[(0)*257 + 1] = w_tile[1]
- Thread 65: w_tile[1*257 + 1] = w_tile[258]
- Thread 66: w_tile[2*257 + 1] = w_tile[515]
Still adjacent writes to consecutive banks. Perfect.
And for the compute phase:
```cpp
for (int k = 0; k < BK; k++) {
float w_val = w_tile[k * BN_PAD + tid];
for (int m = 0; m < M; m++) {
accum[m] += x_tile[m * BK + k] * w_val;
}
}
```
At fixed k, threads access w_tile[k*257 + 0], w_tile[k*257 + 1], ... These are contiguous in shared memory. No bank conflicts.
And for x_tile: stored as x_tile[m * BK + k]. In compute:
```cpp
x_tile[m * BK + k]
```
For fixed k and varying m: addresses m*BK + k are BK apart. With BK=64, that's 256 bytes = 64 banks apart. 64 % 32 = 0. All bank 0! Bank conflicts!
To fix x_tile, store as x_tile[k * (M+pad) + m] or x_tile[k * M + m].
If x_tile[k * M + m], then compute access `x_tile[k * M + m]` is contiguous for varying m. Perfect.
Load loop for x_tile:
```cpp
for (int idx = tid; idx < M * BK; idx += BN) {
int m = idx / BK; // Wait, this was based on BK elements per row
int k = idx % BK;
x_tile[k * M + m] = x[m * K + k_off + k];
}
```
For idx = tid: m = tid/64, k = tid%64.
Thread 0: m=0, k=0 -> x_tile[0] = x[0*K + k_off + 0]
Thread 1: m=0, k=1 -> x_tile[1 * 128 + 0] = x_tile[128] = x[0*K + k_off + 1]
Wait, x_tile[128] and x_tile[0] are 128 floats = 512 bytes apart. 512/4 = 128 banks = 0 mod 32. Same bank! For threads 0 and 1 writing to x_tile[0] and x_tile[128], they hit the same bank.
But actually, if we set the pad to 1: x_tile[k * (M+1) + m] = x_tile[k * 129 + m].
Thread 0: x_tile[0] = x[0]
Thread 1: x_tile[129] = x[1]
129 floats = 516 bytes. 516/4 = 129. 129 % 32 = 1. Thread 1 writes to bank 1. Adjacent threads write to adjacent banks!
Actually, M=128 is a multiple of 32. So M * 4 bytes = 512 bytes = 128 banks. 128 % 32 = 0. So x_tile[k*M + m] for m=0 and m=1 are 512 bytes apart, hitting bank 0 and bank 0 again. Bank conflicts.
With x_tile[k * (M+1) + m]: offset between m=0 and m=1 is 129*4 = 516 bytes. 516/4 = 129. 129 % 32 = 1. So m=0 hits bank 0, m=1 hits bank 1, m=2 hits bank 2, etc. Perfect!
But the access pattern for compute would be:
```cpp
x_tile[k * 129 + m]
```
For fixed k and different m, threads access the same k*129 portion, but different m. Since all threads use the same m in lockstep (iterating m from 0 to 127), all threads in a warp access x_tile[k*129 + m] for the SAME m simultaneously. That's a broadcast! All threads read the same address. No bank conflicts!
Wait, yes! In the compute loop:
```cpp
for (int k = 0; k < BK; k++) {
float w_val = w_tile[k * BN_PAD + tid];
for (int m = 0; m < M; m++) {
accum[m] += x_tile[k * M_PAD + m] * w_val;
}
}
```
At a given step in the m loop, ALL threads access x_tile[k * M_PAD + m] for the same m. This is a broadcast to all threads. Shared memory handles broadcasts efficiently. So we don't even need a pad for x_tile for the compute reads!
But for the write during loading:
```cpp
x_tile[k * M + m] = x[m * K + k_off + k];
```
Threads in a warp write to x_tile for different (m, k) values. With x_tile[k * M + m], threads 0 and 1 access x_tile[0*128+0]=x_tile[0] and x_tile[1*128+0]=x_tile[128] (assuming bn calculation was different, but for x it's m and k).
Wait, for the x loading loop:
idx = tid + i*BN. m = idx / BK. k = idx % BK.
Thread 0: idx=0, m=0, k=0 -> x_tile[0 * M + 0] = x_tile[0]
Thread 1: idx=1, m=0, k=1 -> x_tile[1 * M + 0] = x_tile[128]
Thread 2: idx=2, m=0, k=2 -> x_tile[2 * M + 0] = x_tile[256]
Hmm, all threads with bn=0 write to x_tile[k*M + 0] for different k. These are M*4 = 512 bytes apart. All same bank (0). Bank conflicts.
If we use M_PAD = 129:
Thread 0: x_tile[0 * 129 + 0] = x_tile[0]
Thread 1: x_tile[1 * 129 + 0] = x_tile[129] -> bank 129%32 = 1
Thread 2: x_tile[2 * 129 + 0] = x_tile[258] -> bank 258%32 = 2
Adjacent threads write to adjacent banks! No conflicts!
Great. So set M_PAD = 129 and BN_PAD = 257.
For w_tile compute access:
`w_tile[k * BN_PAD + tid]`
For fixed k and varying tid: contiguous in shared memory, no conflicts.
For w_tile load:
```cpp
for (int idx = tid; idx < BN * BK; idx += BN) {
int bn = idx / BK;
int k = idx % BK;
w_tile[k * BN_PAD + bn] = weight[(n_start + bn) * K + k_off + k];
}
```
Thread 0: idx=0, bn=0, k=0 -> w_tile[0*257 + 0] = w_tile[0]
Thread 1: idx=1, bn=0, k=1 -> w_tile[1*257 + 0] = w_tile[257] -> bank 257%32 = 1
Thread 2: idx=2, bn=0, k=2 -> w_tile[2*257 + 0] = w_tile[514] -> bank 514%32 = 2
Adjacent threads write to adjacent banks. No conflicts.
For x_tile load:
```cpp
for (int idx = tid; idx < M * BK; idx += BN) {
int m = idx / BK;
int k = idx % BK;
x_tile[k * M_PAD + m] = x[m * K + k_off + k];
}
```
Thread 0: m=0, k=0 -> x_tile[0 * 129 + 0] = x_tile[0]
Thread 1: m=0, k=1 -> x_tile[1 * 129 + 0] = x_tile[129] -> bank 129%32 = 1
Thread 2: m=0, k=2 -> x_tile[2 * 129 + 0] = x_tile[258] -> bank 258%32 = 2
Great!
But for x_tile compute access:
`x_tile[k * M_PAD + m]`
For fixed k and fixed m, all threads access the same address (broadcast). No conflicts.
For x_tile, we actually don't need M_PAD for compute since it's broadcast. But for load, we do.
Actually, we could just use a larger M_PAD to avoid bank conflicts during load. With M_PAD = 129, x_tile size = BK * M_PAD = 64 * 129 = 8256 floats = 33KB. Slightly more than M*BK = 8192 floats (32KB).
With BN_PAD = 257, w_tile size = BK * BN_PAD = 64 * 257 = 16448 floats = 64KB + 256 floats. Slightly more than BN*BK = 16384 floats (64KB).
Total shared memory: ~33KB + ~65KB = 98KB. Still under 128KB per block.
Wait, I need to make sure M_PAD and BN_PAD are compile-time constants. Let me set:
```cpp
#define M_PAD 129
#define BN_PAD 257
```
Or to be more precise, M_PAD should be the next multiple of 32 plus 1? Actually, the bank size on H100 is 4 bytes, and there are 32 banks. To avoid conflicts when accessing elements spaced by M_PAD floats, we need M_PAD * 4 % (32 * 4) != 0, i.e., M_PAD % 32 != 0. Since M=128, M%32=0. So M_PAD = M+1 = 129 works. 129 % 32 = 1. Good.
Similarly, BN_PAD = BN + 1 = 257. 257 % 32 = 1. Good.
We could also use M+8 = 136, which is a nicer number and 136 % 32 = 8. Let me check: 136 * 4 = 544 bytes. 544 / 128 bytes per bank-group? Actually 32 banks * 4 bytes = 128 bytes per full cycle. 544 % 128 = 544 - 4*128 = 544 - 512 = 32. 32/4 = 8. So it would hit bank 8 for m=1. Adjacent m values hit adjacent banks. This works too and keeps alignment. Let me use M_PAD = 136 and BN_PAD = 264. That way k*136 is always aligned to something nice.
Wait, x_tile size with M_PAD=136: BK * 136 = 64 * 136 = 8704 floats.
w_tile size with BN_PAD=264: BK * 264 = 64 * 264 = 16896 floats.
Total: 8704 + 16896 = 25600 floats = 102.4KB. Still under 128KB.
Hmm, that's a lot more shared memory but still OK. Actually, to minimize shared memory, let me use M+1=129 and BN+1=257.
x_tile: 64 * 129 = 8256 floats = 33KB
w_tile: 64 * 257 = 16448 floats = 64KB + 1KB
Total: ~98KB
OK let me go with M_PAD=129 and BN_PAD=257. Though I wonder if I can just use M_PAD=129 and store w_tile without padding since w_tile access is contiguous anyway. Actually for w_tile compute: `w_tile[k * BN_PAD + tid]`. If BN_PAD=BN=256, then consecutive tids access consecutive elements. No bank conflicts needed!
Wait, for w_tile write during loading:
`w_tile[k * BN + bn] = ...`
With BN=256, thread 0 writes w_tile[0], thread 1 writes w_tile[256], thread 2 writes w_tile[512], etc. These are spaced by 256 floats = 1024 bytes. 1024 / 4 = 256 bank positions. 256 % 32 = 0. All bank 0! Conflicts.
So we DO need BN_PAD for w_tile writes.
Similarly for x_tile writes.
OK, let me write the kernel with these pads.
Actually, another option: instead of padding, transpose the access pattern. During loading, threads can cooperate differently.
For w_tile: instead of `for (int idx = tid; idx < BN * BK; idx += BN)`, we can have each warp load a full row of w. But that requires warp-level programming. The padding approach with M_PAD and BN_PAD is simpler.
Let me write the final kernel.
One more thing: the inner compute loop with `x_tile[k * M_PAD + m]`. This is strided by M_PAD = 129. The inner loop iterates m from 0 to 127. So it accesses x_tile[k*129 + 0], x_tile[k*129 + 1], ... These are adjacent in shared memory. Since it's in a single thread, it's just sequential. No bank conflict issues for single thread. And since all threads access the same m (broadcast), also fine.
Now let me write the Python module with the kernel.
Actually, I should also consider using `tanh` for sigmoid? `sigmoid(x) = 0.5 * (1 + tanh(x/2))`. Not sure if faster. `expf` is usually fast enough.
Let me also check: do I need to initialize accumulators to 0? Yes. And do I need to initialize output to 0? Yes.
Here's the final structure of my kernel code:
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#define M 128
#define K 32768
#define N 32768
#define BK 64
#define BN 256
#define M_PAD 129
#define BN_PAD 257
#define NUM_K_TILES (K / BK)
#define NUM_BLOCKS (N / BN)
__global__ void fused_linear_sigmoid_sum_kernel(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output
) {
__shared__ float x_tile[BK * M_PAD];
__shared__ float w_tile[BK * BN_PAD];
__shared__ float warp_sums[(BN / 32) * M];
int n_start = blockIdx.x * BN;
int tid = threadIdx.x;
int n = n_start + tid;
int lane_id = tid % 32;
int warp_id = tid / 32;
int num_warps = BN / 32;
float accum[M];
#pragma unroll
for (int m = 0; m < M; m++) accum[m] = 0.0f;
for (int kt = 0; kt < NUM_K_TILES; kt++) {
int k_off = kt * BK;
// Load x_tile into shared memory
for (int idx = tid; idx < M * BK; idx += BN) {
int m = idx / BK;
int k = idx % BK;
x_tile[k * M_PAD + m] = x[m * K + k_off + k];
}
// Load w_tile into shared memory
for (int idx = tid; idx < BN * BK; idx += BN) {
int bn = idx / BK;
int k = idx % BK;
w_tile[k * BN_PAD + bn] = weight[(n_start + bn) * K + k_off + k];
}
__syncthreads();
// Compute partial dot products
#pragma unroll
for (int k = 0; k < BK; k++) {
float w_val = w_tile[k * BN_PAD + tid];
#pragma unroll
for (int m = 0; m < M; m++) {
accum[m] += x_tile[k * M_PAD + m] * w_val;
}
}
__syncthreads();
}
// Apply sigmoid
float thread_sums[M];
#pragma unroll
for (int m = 0; m < M; m++) {
float val = accum[m] + bias[n];
thread_sums[m] = 1.0f / (1.0f + expf(-val));
}
// Warp-level reduction
#pragma unroll
for (int m = 0; m < M; m++) {
float val = thread_sums[m];
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (lane_id == 0) {
warp_sums[warp_id * M + m] = val;
}
}
__syncthreads();
// Inter-warp reduction and atomic add
if (tid == 0) {
#pragma unroll
for (int m = 0; m < M; m++) {
float val = 0.0f;
#pragma unroll
for (int w = 0; w < num_warps; w++) {
val += warp_sums[w * M + m];
}
atomicAdd(&output[m], val);
}
}
}
'''
import torch
import torch.nn as nn
class ModelNew(nn.Module):
def __init__(self, input_size, hidden_size):
super(ModelNew, self).__init__()
self.linear = nn.Linear(input_size, hidden_size)
def forward(self, x):
# If sizes don't match our optimized kernel, fall back
# But they always will for this task
if x.shape[0] == M and x.shape[1] == K and self.linear.weight.shape[0] == N and self.linear.weight.shape[1] == K:
output = torch.zeros(M, 1, device=x.device, dtype=torch.float32)
# ... call kernel ...
return output
else:
x = self.linear(x)
x = torch.sigmoid(x)
x = torch.sum(x, dim=1, keepdim=True)
return x
```
Wait, I can't use M, K, N in Python since they're C preprocessor macros. I'll check the shapes directly.
Actually, for the `cpp_sources`, I need to include the kernel and a wrapper function.
Also, in the wrapper function, I should check that the tensors are contiguous and on CUDA.
Let me write the complete module:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#define M 128
#define K 32768
#define N 32768
#define BK 64
#define BN 256
#define M_PAD 129
#define BN_PAD 257
#define NUM_K_TILES (K / BK)
#define NUM_BLOCKS (N / BN)
__global__ void fused_linear_sigmoid_sum_kernel(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output
) {
__shared__ float x_tile[BK * M_PAD];
__shared__ float w_tile[BK * BN_PAD];
__shared__ float warp_sums[(BN / 32) * M];
int n_start = blockIdx.x * BN;
int tid = threadIdx.x;
int n = n_start + tid;
int lane_id = tid % 32;
int warp_id = tid / 32;
const int num_warps = BN / 32;
float accum[M];
#pragma unroll
for (int m = 0; m < M; m++) accum[m] = 0.0f;
for (int kt = 0; kt < NUM_K_TILES; kt++) {
int k_off = kt * BK;
for (int idx = tid; idx < M * BK; idx += BN) {
int m = idx / BK;
int k = idx % BK;
x_tile[k * M_PAD + m] = x[m * K + k_off + k];
}
for (int idx = tid; idx < BN * BK; idx += BN) {
int bn = idx / BK;
int k = idx % BK;
w_tile[k * BN_PAD + bn] = weight[(n_start + bn) * K + k_off + k];
}
__syncthreads();
#pragma unroll
for (int k = 0; k < BK; k++) {
float w_val = w_tile[k * BN_PAD + tid];
#pragma unroll
for (int m = 0; m < M; m++) {
accum[m] += x_tile[k * M_PAD + m] * w_val;
}
}
__syncthreads();
}
float thread_sums[M];
#pragma unroll
for (int m = 0; m < M; m++) {
float val = accum[m] + bias[n];
thread_sums[m] = 1.0f / (1.0f + expf(-val));
}
#pragma unroll
for (int m = 0; m < M; m++) {
float val = thread_sums[m];
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (lane_id == 0) {
warp_sums[warp_id * M + m] = val;
}
}
__syncthreads();
if (tid == 0) {
#pragma unroll
for (int m = 0; m < M; m++) {
float val = 0.0f;
#pragma unroll
for (int w = 0; w < num_warps; w++) {
val += warp_sums[w * M + m];
}
atomicAdd(&output[m], val);
}
}
}
void fused_linear_sigmoid_sum_cuda(
torch::Tensor x,
torch::Tensor weight,
torch::Tensor bias,
torch::Tensor output
) {
const dim3 blocks(NUM_BLOCKS);
const dim3 threads(BN);
fused_linear_sigmoid_sum_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>()
);
}
'''.replace('M 128', 'M 128').replace('K 32768', 'K 32768').replace('N 32768', 'N 32768')
# Actually, for flexibility, we might want to not hardcode and instead use runtime checks
# But the problem says specific sizes
```
Wait, in the C++ wrapper, I need to include ATen headers. But `load_inline` with `cpp_sources` that use `torch::Tensor` should work. I need to add `#include <torch/extension.h>` in the header.
Let me write the complete inline source properly.
Actually, the `functions` parameter is for the host wrapper functions. `load_inline` generates the pybind11 bindings for those.
One issue: our kernel uses `atomicAdd` on `output[m]`. But `output` is a tensor of shape (M, 1) or (M,). In our case, we allocate it as (M,) but the model's output should be (M, 1). So in Python, I'll allocate `output = torch.zeros(M, device=x.device, dtype=torch.float32)` and then reshape to (M, 1) after the kernel.
Actually, it's easier to allocate output as `torch.zeros(batch_size, 1, device=x.device)` and pass `.view(-1)` to the kernel, or allocate as `.squeeze()`. Let me just allocate `output = torch.zeros(batch_size, device=x.device, dtype=torch.float32)` and then return `output.unsqueeze(1)`.
Let me now put together the full kernel code.
I should also consider whether to template the kernel or just hardcode. Hardcoding is simpler and might be faster due to better unrolling. But we should verify with compile and correctness.
Wait, one more thing. The compute loop:
```cpp
for (int k = 0; k < BK; k++) {
float w_val = w_tile[k * BN_PAD + tid];
for (int m = 0; m < M; m++) {
accum[m] += x_tile[k * M_PAD + m] * w_val;
}
}
```
Since x_tile is accessed as `x_tile[k * M_PAD + m]`, for fixed k this is a contiguous segment. The compiler should be able to unroll this nicely.
However, `accum[m]` is in registers. Each iteration adds to a different register. The FMA instructions should be independent because they target different accumulators. The compiler should interleave them nicely.
But there's a potential issue: `w_tile[k * BN_PAD + tid]` is loaded once and then broadcast to all M FMAs. The compiler might keep `w_val` in a register, which is fine.
One optimization I could try later is processing M in groups to improve instruction-level parallelism. But for now, let's see if this compiles and passes correctness.
Write the code, compile, and test!