reasoning
Let me analyze the reference model:
1. `nn.Linear(input_size, hidden_size)` - matrix multiplication: (batch_size, input_size) @ (hidden_size, input_size).T = (batch_size, hidden_size)
2. `x = x * scale_factor` - scale by 2.0
3. `x = x + x` - add residual (which is just 2x, so combined with step 2 it's 4x)
4. `torch.clamp(x, clamp_min, clamp_max)`
5. `torch.logsumexp(x, dim=1, keepdim=True)` - reduces hidden_size dim to 1
6. `x = x * torch.nn.functional.mish(x)` - mish activation then multiply
Let me compute the total operations:
- Matmul: batch_size * input_size * hidden_size * 2 = 1024 * 8192 * 8192 * 2 ≈ 137 GFLOPs
- The rest is memory-bound, operating on batch_size * hidden_size elements
The output shape is (batch_size, 1) = (1024, 1).
One key observation: `x = x + x` is just doubling x again. Combined with `x = x * scale_factor` where scale_factor=2.0, the total scaling is 4x.
So mathematically:
- Start with matmul output M
- x = M * 2.0
- x = x + x = 2x = 4M
- clamp(x, -10, 10)
- logsumexp(clamped, dim=1, keepdim=True)
- result = LSE * mish(LSE)
I can fuse most of this into a single kernel. The main challenge is the matmul (137 GFLOPs). For H100, using cuBLAS for the matmul would be ideal, but I need custom CUDA for the fusion.
Wait, the sizes are 1024 x 8192 x 8192. This is actually a very large matmul. For H100 at 989 TFLOPS (TF32 Tensor Core), 137 GFLOPs would take ~0.14 ms. Memory bandwidth is 3.35 TB/s.
For the rest:
- After matmul: 1024 * 8192 = 8.4M elements
- clamp + scaling: trivial, memory-bound
- logsumexp: needs to read 8.4M elements, write 1024 elements
- final mish + multiply: read 1024, write 1024
So the matmul dominates. Let me think about what's the best approach:
1. Use cuBLAS/lt for matmul (optimal for this size)
2. Fuse the rest into custom kernels
Actually, is there a way to fuse some operations with the matmul? The scaling by 4 and clamp are element-wise, so they can be fused. But logsumexp is a reduction across hidden_size.
An efficient approach:
1. cuBLAS matmul: computes M = x @ W.T
2. Custom kernel that does: compute 4*M, clamp, then compute logsumexp per row, store clamped values AND LSE result
3. Actually, logsumexp doesn't need the clamped values stored - we can compute it directly.
Wait, let me re-read the code more carefully:
```python
x = self.matmul(x) # (batch_size, hidden_size)
x = x * self.scale_factor # scale
x = x + x # double again
x = torch.clamp(x, self.clamp_min, self.clamp_max) # clamp
x = torch.logsumexp(x, dim=1, keepdim=True) # LSE over hidden_size -> (batch_size, 1)
x = x * torch.nn.functional.mish(x) # mish activation then multiply
return x # (batch_size, 1)
```
The key insight: after logsumexp, the output is (batch_size, 1). The intermediate tensor of size (batch_size, hidden_size) doesn't need to be stored in full between operations.
So the optimal strategy is:
1. Use cuBLAS for matmul: M = x @ W.T (output shape: batch_size x hidden_size)
2. In a custom kernel, for each row, read all hidden_size elements, apply 4x scale, clamp, compute logsumexp, output 1 value per row
3. Another tiny kernel for mish(x) * x on the batch_size elements
Actually, I can fuse step 2 and 3 too. The kernel can:
- For each row: compute scaled & clamped values, compute logsumexp, write result to output
- Then another kernel or the same if I use multiple passes... but actually for step 3 it's a simple element-wise on batch_size elements.
Let me think about the logsumexp computation. For numerical stability, LSE is typically computed as:
```python
max_val = max(x)
lse = max_val + log(sum(exp(x - max_val)))
```
For the fused kernel per row:
- First pass over the row: find max value
- Second pass: compute sum of exp(x - max)
- LSE = max + log(sum)
This requires 2 reads of the row, or I can cache the row in shared memory. With hidden_size=8192 and batch_size=1024:
- Each row has 8192 floats = 32KB
- Shared memory per SM is 228KB, so we can fit ~7 rows fully.
- But we can also process each row in tiles.
For H100, I should use warp-level or block-level reductions.
Let me design the kernel architecture:
Option A: One block per row
- Block size: maybe 1024 threads
- Each thread processes 8192/1024 = 8 elements
- Phase 1: load elements, scale, clamp, find max using warp reduce then block reduce
- Phase 2: load elements again (or keep in registers), compute exp(x - max), sum using warp reduce then block reduce
- Phase 3: compute LSE = max + log(sum)
- Write LSE for this row
This reads the matmul output twice. Alternatively, store in shared memory:
- Shared memory: 8192 floats = 32KB. With 1024 threads, that's manageable.
- We can store the entire row in shared memory after scaling and clamping.
Actually, for performance, let's think about whether using shared memory is better. The main memory access is reading the matmul result (32KB per row). Writing is 4 bytes per row. So it's very read-heavy.
One block per row with 1024 threads:
- Each thread handles 8 consecutive elements (good coalescing)
- First pass: load 8 elements, clamp, find local max
- Block reduce for global max
- Second pass: reload 8 elements (or reuse from registers?), compute exp(x - max)
- But wait, I loaded them into registers in the first pass. Can I reuse? No, because I need the global max first.
- So I either: (1) load twice, or (2) store to shared memory in first pass and read from there in second pass.
For 8192 elements:
- Option 1: 2 reads from global memory, 0 shared memory for data = 64KB reads per row
- Option 2: 1 read from global memory, 1 write + 1 read from shared memory = 32KB global + 32KB shared
- Shared memory is much faster than global memory
But actually, if I use 1024 threads and each loads 8 elements:
- First pass: coalesced read of 32 floats per warp (128 bytes), perfectly aligned
- Write to shared memory: also coalesced
- Second pass: read from shared memory, also fast
Shared memory approach seems better: it reduces global memory reads by half.
Let me also think about whether to fuse the final mish * mish computation. After LSE, we have 1024 values. Then:
```python
x = x * mish(x)
```
where mish(x) = x * tanh(softplus(x)) = x * tanh(log(1 + exp(x)))
This is a simple element-wise kernel on 1024 elements.
Now, for the matmul: I should use `torch.matmul` or `F.linear` which will call cuBLAS. This is optimal for 1024x8192x8192.
Total kernel design:
1. In `__init__`: Create weight and bias (nn.Linear handles this, but I can store them directly)
2. In `forward`:
- y = F.linear(x, self.weight, self.bias) # or x @ W.T + b
- Call custom CUDA kernel to: scale by 4, clamp, compute LSE per row
- Call custom CUDA kernel for mish * x
Wait, the bias is part of nn.Linear. Looking at the reference:
```python
self.matmul = nn.Linear(input_size, hidden_size)
```
So there is a bias. I should create a weight and bias in ModelNew, or just call the linear operation.
I need to handle the initialization correctly. The constructor receives `[input_size, hidden_size, scale_factor, clamp_min, clamp_max]`.
Let me write the kernel:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
cpp_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>
#define BLOCK_SIZE 1024
__inline__ __device__ float warp_reduce_max(float val) {
for (int offset = 16; offset > 0; offset /= 2)
val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
return val;
}
__inline__ __device__ float warp_reduce_sum(float val) {
for (int offset = 16; offset > 0; offset /= 2)
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
return val;
}
__global__ void fused_clamp_lse_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int hidden_size,
float scale,
float clamp_min,
float clamp_max
) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_ptr = input + row * hidden_size;
// Stage 1: Find max after scaling and clamping
float local_max = -INFINITY;
for (int idx = threadIdx.x; idx < hidden_size; idx += BLOCK_SIZE) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
local_max = fmaxf(local_max, val);
}
local_max = warp_reduce_max(local_max);
__shared__ float s_max;
if ((threadIdx.x & 31) == 0) {
atomicMax((int*)&s_max, __float_as_int(local_max));
}
__syncthreads();
// Note: atomicMax with float bits doesn't work for negative values properly
// Better: use shared memory reduction
}
```
Actually, atomicMax on float bits is problematic for negative values. Let me use a proper block reduction.
Let me redesign the reduction:
```cpp
__inline__ __device__ float block_reduce_max(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_max(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : -INFINITY;
if (wid == 0) val = warp_reduce_max(val);
return val;
}
__inline__ __device__ float block_reduce_sum(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
if (wid == 0) val = warp_reduce_sum(val);
return val;
}
```
Wait, `blockDim.x / 32` = 1024/32 = 32, which fits in shared[32].
Now the kernel:
```cpp
__global__ void fused_clamp_lse_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int hidden_size,
float scale,
float clamp_min,
float clamp_max
) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_ptr = input + row * hidden_size;
extern __shared__ float sdata[];
// Stage 1: Load, scale, clamp to shared memory and find max
float local_max = -INFINITY;
for (int idx = threadIdx.x; idx < hidden_size; idx += BLOCK_SIZE) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
sdata[idx] = val; // Store to shared memory
local_max = fmaxf(local_max, val);
}
__syncthreads();
float row_max = block_reduce_max(local_max);
// Broadcast max to all threads
if (threadIdx.x == 0) {
sdata[0] = row_max;
}
__syncthreads();
row_max = sdata[0];
// Stage 2: Compute sum of exp(x - max)
float local_sum = 0.0f;
for (int idx = threadIdx.x; idx < hidden_size; idx += BLOCK_SIZE) {
local_sum += expf(sdata[idx] - row_max);
}
float row_sum = block_reduce_sum(local_sum);
if (threadIdx.x == 0) {
output[row] = row_max + logf(row_sum);
}
}
```
Wait, `sdata[idx] = val` with idx going from threadIdx.x to hidden_size in steps of BLOCK_SIZE. This means thread 0 writes to sdata[0], sdata[1024], sdata[2048], etc. That's not coalesced for shared memory! Shared memory works best with consecutive threads accessing consecutive addresses.
Actually, for shared memory, there's no coalescing requirement like global memory, but bank conflicts matter. If thread 0 accesses bank 0, thread 1 accesses bank 1, etc., there's no conflict. But in my loop, thread 0 accesses 0, 1024, 2048, etc. - all in the same bank (bank 0 if 1024 is divisible by 32, which it is). So there would be massive bank conflicts!
Let me fix this. Better to have each thread handle consecutive elements:
```cpp
int tid = threadIdx.x;
int num_elements_per_thread = (hidden_size + BLOCK_SIZE - 1) / BLOCK_SIZE;
int start = tid * num_elements_per_thread;
int end = min(start + num_elements_per_thread, hidden_size);
for (int idx = start; idx < end; idx++) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
sdata[idx] = val;
local_max = fmaxf(local_max, val);
}
```
But now the global memory reads are not perfectly coalesced (consecutive threads read consecutive elements, which is coalesced! thread 0 reads 0..7, thread 1 reads 8..15, etc. So global reads are coalesced. Shared writes are also consecutive per thread but strided across threads - there could be bank conflicts if multiple threads try to write at the same time... no, they write to different banks since consecutive addresses are in consecutive banks (each bank is 4 bytes, 32 banks). Thread 0 writes to bank 0, thread 1 writes to bank 1, etc. No bank conflicts as long as each thread writes to consecutive addresses.
Actually wait, if thread 0 writes sdata[0], sdata[1], ... sdata[7], these happen sequentially within the thread, not simultaneously with other threads. There's no bank conflict in time. So this is fine.
But for the second pass, if all threads read sdata in a similar pattern, same thing - no bank conflicts.
However, for the second pass, if I do:
```cpp
for (int idx = start; idx < end; idx++) {
local_sum += expf(sdata[idx] - row_max);
}
```
This is fine.
But wait, for very large hidden_size (8192) and BLOCK_SIZE=1024, each thread handles 8 elements. That's fine.
Actually, let me reconsider whether storing in shared memory is worth it. Shared memory size per block is limited. For 8192 floats = 32KB. With 1024 threads per block and 32 blocks on an SM (max), that's 1MB, which is under 228KB per SM... wait, 32 blocks * 32KB = 1MB > 228KB. So we can't have 32 blocks concurrently.
But we don't need 32 blocks per SM. With 1024 threads per block, on H100, each block uses at least 32 warps. Occupancy might not be great, but this is a reduction kernel. Let me think...
Actually, 1024 threads per block with 32KB shared memory:
- Register usage: probably ~32 registers per thread = 32K registers per block
- Shared memory: 32KB per block
- On H100: 64K registers per SM, 228KB shared memory per SM
- So at most 2 blocks per SM (register limited, since 2 * 32K = 64K registers). Actually 1024 threads * 32 registers = 32,768 registers. H100 has 64K registers per SM. So 2 blocks per SM max.
- But with 100% occupancy of 2 blocks = 2048 threads per SM. H100 can run up to 2048 threads per SM. So actually 2 blocks per SM gives 100% occupancy? Wait, max threads per SM is 2048. Yes, 2 blocks of 1024 = 2048 threads = 100% occupancy if register limited.
Let me think about whether 1024 threads per block is optimal. Maybe 512 threads with 2 elements each?
Actually, let me stick with 1024 threads per block.
For the shared memory size: I declared `extern __shared__ float sdata[];` which means dynamic shared memory. I need to pass the size in the kernel launch.
Wait, `sdata[ idx ]` with idx up to 8191 and extern shared memory. I need `hidden_size * sizeof(float)` bytes. Let me use dynamic shared memory.
Kernel launch:
```cpp
fused_clamp_lse_kernel<<<batch_size, BLOCK_SIZE, hidden_size * sizeof(float)>>>(...)
```
Hmm, 32KB per block. With 1024 blocks total (batch_size=1024), this will run. But can we optimize further?
Actually, I just realized: do I even need to store all values in shared memory? I can do the logsumexp in a single pass if I use the streaming log-sum-exp algorithm:
```python
m = -inf
s = 0
for x in values:
if x > m:
s = s * exp(m - x) + 1
m = x
else:
s = s + exp(x - m)
lse = m + log(s)
```
This is a sequential algorithm but can be parallelized with tree reduction. However, for per-thread iteration within a warp, it requires sequential processing within each thread, then warp-level reduction using the same formula. This is more complex but avoids the second pass (or shared memory).
Let me think if this is worth it. The two-pass approach with shared memory reads the global memory once and does two passes over shared memory. The single-pass approach reads global memory once and does a single pass. But the single-pass parallel reduction is trickier.
Actually, the two-pass approach with shared memory is:
1. Read global memory once
2. Write shared memory once
3. Read shared memory once
4. Do two reductions (max and sum)
Total per element: 1 global read, 1 shared write, 1 shared read, plus the extra operations.
With 8192 elements and 1024 threads, this should be very fast. The bottleneck might be the global memory read, and we only do one read.
Wait, I need to make sure the first pass writes to shared memory correctly. Let me think about occupancy.
Actually, there's another issue. The LSE computation in PyTorch also uses numerical stabilization. My two-pass approach with `exp(x - max)` is standard and numerically stable as long as `max` is computed correctly.
But what if all values are -inf or the row is empty? Not possible here since input is from matmul of random positive values.
Let me reconsider whether we need shared memory at all. If I just do two passes:
Pass 1 (find max):
```cpp
float local_max = -INFINITY;
for (int idx = start; idx < end; idx++) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
local_max = fmaxf(local_max, val);
}
float row_max = block_reduce_max(local_max);
__syncthreads(); // to broadcast
```
Pass 2 (compute sum):
```cpp
float local_sum = 0.0f;
for (int idx = start; idx < end; idx++) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
local_sum += expf(val - row_max);
}
float row_sum = block_reduce_sum(local_sum);
```
This reads global memory twice but doesn't need shared memory. With 64KB per row and 1024 rows = 64MB total reads. Memory bandwidth is 3.35 TB/s, so 64MB takes ~0.02 ms. The matmul takes ~0.14 ms. So the overhead of reading global memory twice vs once is small.
If I use shared memory, I save ~0.01 ms but might have lower occupancy due to shared memory usage.
Actually, for 32KB shared memory per block and 228KB per SM, we can fit 7 blocks per SM (shared memory). But register limit gives us 2 blocks. So shared memory isn't the limiting factor for occupancy - registers are.
Hmm, let me just use the simpler two-pass global read approach. Lower complexity, fewer bugs, and the performance difference might be negligible.
Wait, one optimization: I could precompute the scale/clamp values and reuse them in the second pass... no, each thread handles the same elements in both passes, so it would need to store 8 values in registers. That's fine! 8 * 4 bytes = 32 bytes per thread, negligible.
```cpp
// Each thread loads its elements into registers
float vals[8]; // for 8192/1024 = 8 elements per thread
int start = threadIdx.x * 8;
int end = min(start + 8, hidden_size);
for (int i = 0; i < 8 && start + i < hidden_size; i++) {
float val = row_ptr[start + i] * scale;
vals[i] = fminf(fmaxf(val, clamp_min), clamp_max);
local_max = fmaxf(local_max, vals[i]);
}
// ... reduce to find row_max ...
float local_sum = 0.0f;
for (int i = 0; i < 8 && start + i < hidden_size; i++) {
local_sum += expf(vals[i] - row_max);
}
```
This only reads global memory once, stores in registers (or recomputes), and doesn't need shared memory. But compute is cheap, global reads are expensive. So this is good.
However, 8192/1024 = 8 exactly. So no bounds checking needed if we ensure exact division.
Wait, hidden_size could be anything. Better keep it general. For this problem it's 8192, but to be safe.
Actually, I can use unrolled loops or just a simple loop. With 8 iterations, the compiler will unroll if we're lucky.
But there's another consideration: what if hidden_size is not divisible by block_size? Let me keep it general with a loop.
Wait, I said start = threadIdx.x * num_elements_per_thread. If hidden_size = 8192 and block_size = 1024, then num_elements_per_thread = 8. start = tid * 8, end = start + 8. This works perfectly.
But for coalescing: thread 0 accesses 0, 1, 2, ..., 7; thread 1 accesses 8, 9, ..., 15. Global memory reads from consecutive threads are accessing consecutive memory, which is perfect coalescing! However, each thread does 8 scalar loads. If the compiler uses vectorized loads (LDG.128), that would be 2 128-bit loads per thread. But since it's consecutive within the thread and consecutive across threads... actually for LDG.128, the warp accesses 128 bytes per instruction, with each thread loading 16 bytes (4 floats). Wait, LDG.128 loads 128 bits = 16 bytes = 4 floats per thread.
For perfect vectorization: warp reads 128 bytes = 32 floats. With 32 threads each doing LDG.128, that's 128 floats per warp instruction? No, LDG.128 means each thread loads 128 bits = 16 bytes. 32 threads * 16 bytes = 512 bytes = 128 floats per warp instruction. But our elements are 4 bytes each.
Actually, for 8192 floats = 32KB per row. With block_size 1024, the warp reads 32 floats (128 bytes) per instruction at a time if doing 4-byte loads. For a row, the accesses are perfectly contiguous.
Let me not overthink the vectorization and just write clean code. The compiler is pretty good.
Here's my revised kernel plan:
```cpp
__global__ void fused_scale_clamp_lse_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int hidden_size,
float scale,
float clamp_min,
float clamp_max
) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_ptr = input + row * hidden_size;
int tid = threadIdx.x;
int num_threads = blockDim.x;
// Phase 1: find max of scaled and clamped values
float local_max = -INFINITY;
for (int idx = tid; idx < hidden_size; idx += num_threads) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
local_max = fmaxf(local_max, val);
}
// Warp reduce
local_max = warp_reduce_max(local_max);
// Block reduce
__shared__ float s_max;
if ((tid & 31) == 0) {
// Use atomic for max
atomicMax((int*)&s_max, __float_as_int(local_max));
}
__syncthreads();
float row_max = s_max;
// Phase 2: compute sum of exp(x - max)
float local_sum = 0.0f;
for (int idx = tid; idx < hidden_size; idx += num_threads) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
local_sum += expf(val - row_max);
}
// Warp reduce
local_sum = warp_reduce_sum(local_sum);
// Block reduce
__shared__ float s_sum;
if ((tid & 31) == 0) {
atomicAdd(&s_sum, local_sum);
}
__syncthreads();
if (tid == 0) {
output[row] = row_max + logf(s_sum);
}
}
```
Wait, the atomicMax for float is problematic. Let me use the proper block reduce approach with shared memory.
Actually, let me use a simpler approach:
```cpp
__device__ float block_reduce_max(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_max(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -INFINITY;
if (wid == 0) val = warp_reduce_max(val);
return val;
}
__device__ float block_reduce_sum(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
if (wid == 0) val = warp_reduce_sum(val);
return val;
}
```
Then the kernel becomes cleaner:
```cpp
float local_max = -INFINITY;
for (int idx = tid; idx < hidden_size; idx += blockDim.x) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
local_max = fmaxf(local_max, val);
}
float row_max = block_reduce_max(local_max);
__syncthreads(); // need this before second phase to avoid shared memory conflicts
float local_sum = 0.0f;
for (int idx = tid; idx < hidden_size; idx += blockDim.x) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
local_sum += expf(val - row_max);
}
float row_sum = block_reduce_sum(local_sum);
if (threadIdx.x == 0) {
output[row] = row_max + logf(row_sum);
}
```
This reads global memory twice but is clean. For performance, I could store values in shared memory or registers, but given the matmul dominates, this should be fine.
Wait, I can actually avoid the second global read by caching values in shared memory. Let me think if it's worth it.
32KB shared memory for storing the row. This might reduce occupancy from 2 blocks to 2 blocks (register limited anyway). So let's use shared memory to save global bandwidth.
```cpp
__global__ void fused_scale_clamp_lse_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int hidden_size,
float scale,
float clamp_min,
float clamp_max
) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_ptr = input + row * hidden_size;
int tid = threadIdx.x;
extern __shared__ float sdata[];
// Load, scale, clamp, store to shared memory
float local_max = -INFINITY;
for (int idx = tid; idx < hidden_size; idx += blockDim.x) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
sdata[idx] = val;
local_max = fmaxf(local_max, val);
}
__syncthreads();
float row_max = block_reduce_max(local_max);
__syncthreads();
// Compute sum of exp(x - max)
float local_sum = 0.0f;
for (int idx = tid; idx < hidden_size; idx += blockDim.x) {
local_sum += expf(sdata[idx] - row_max);
}
float row_sum = block_reduce_sum(local_sum);
if (threadIdx.x == 0) {
output[row] = row_max + logf(row_sum);
}
}
```
Shared memory allocation: `hidden_size * sizeof(float)`.
One concern: bank conflicts. If `hidden_size = 8192` and blockDim.x = 1024, then in the first loop:
- Thread 0 writes sdata[0], sdata[1024], sdata[2048], ...
- sdata[0] is bank 0, sdata[1024] is bank 0 (1024 / 32 = 32, so 1024 % 32 = 0). Bank 0!
- Thread 1 writes sdata[1], sdata[1025], ... sdata[1025] is bank 1.
- But thread 0 and thread 1 write simultaneously. Thread 0 writes to bank 0, thread 1 to bank 1, no conflict.
- However, thread 0 writes sdata[0] and sdata[1024] sequentially, not simultaneously with other threads writing to bank 0.
Actually, in a loop, thread 0 writes sdata[0] in iteration 0, then in iteration 1 it writes sdata[1024]. Other threads write other banks simultaneously. The issue is: when thread 0 writes sdata[1024] (bank 0), which other thread is also writing to bank 0? Thread 32 writes sdata[1056] = sdata[1024 + 32], which is bank 0. But they are in the same warp! Threads 0 and 32 are in the same warp. Both try to write to bank 0 simultaneously. This is a 2-way bank conflict.
Actually, let me check: bank = (address / 4) % 32.
- sdata[1024]: address = 4096, bank = (1024) % 32 = 0
- sdata[1056]: address = 4224, bank = (1056) % 32 = 0
Yes, both threads 0 and 32 write to bank 0 simultaneously. In fact, threads 0, 32, 64, ... all write to bank 0 in the second iteration. That's a 32-way bank conflict! This is terrible.
I need to fix the shared memory bank conflicts.
Option 1: Reorder the loop so threads write to consecutive addresses. But this changes the coalescing of global memory reads.
Actually, if I write:
```cpp
int elements_per_thread = (hidden_size + blockDim.x - 1) / blockDim.x;
int start = tid * elements_per_thread;
int end = min(start + elements_per_thread, hidden_size);
for (int idx = start; idx < end; idx++) {
...
sdata[idx] = val;
}
```
Then:
- Thread 0 writes sdata[0], sdata[1], ..., sdata[7]
- Thread 1 writes sdata[8], sdata[9], ..., sdata[15]
- Thread 0 writes to banks 0, 1, 2, ... sequentially over time
- Thread 1 writes to banks 8, 9, ... sequentially over time
- At any given cycle (within the warp), thread 0 writes to bank X, thread 1 writes to bank X+8, etc. No conflicts!
For global memory reads: thread 0 reads row_ptr[0..7], thread 1 reads row_ptr[8..15]. Scalar loads within a thread might not be vectorized, but the warp as a whole reads consecutive memory. However, since each thread does multiple scalar loads, there could be multiple warp-wide coalesced memory transactions.
Actually, with 8 scalar loads per thread and 32 threads per warp, the compiler might or might not vectorize. But the accesses are still coalesced.
Let me use this approach:
```cpp
int elements_per_thread = (hidden_size + blockDim.x - 1) / blockDim.x;
int start = tid * elements_per_thread;
int end = min(start + elements_per_thread, hidden_size);
for (int idx = start; idx < end; idx++) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
sdata[idx] = val;
local_max = fmaxf(local_max, val);
}
```
Actually wait, for 8192 / 1024 = 8 exactly. So elements_per_thread = 8, start = tid * 8, end = start + 8.
For the second loop:
```cpp
for (int idx = start; idx < end; idx++) {
local_sum += expf(sdata[idx] - row_max);
}
```
No bank conflicts in either phase.
Now, for the final mish kernel:
```cpp
__device__ inline float mish(float x) {
return x * tanhf(logf(1.0f + expf(x)));
}
```
Wait, `log(1 + exp(x))` is softplus. For large x, this is just x. For numerical stability, better to use:
```cpp
float softplus = x > 20.0f ? x : logf(1.0f + expf(x));
return x * tanhf(softplus);
```
Actually, let's look at what PyTorch's mish does. PyTorch mish is defined as `x * tanh(softplus(x))` where softplus uses the stable implementation.
For small values: softplus(x) = log(1 + exp(x))
For large positive x: softplus(x) ≈ x
For large negative x: softplus(x) ≈ exp(x)
A numerically stable implementation:
```cpp
float softplus(float x) {
return x > 0.0f ? x + log1pf(expf(-x)) : log1pf(expf(x));
}
```
And tanh is just `tanhf` in CUDA.
```cpp
__device__ inline float mish(float x) {
float sp = x > 0.0f ? x + log1pf(expf(-x)) : log1pf(expf(x));
return x * tanhf(sp);
}
```
But actually, since we've computed LSE which are values around log(8192 * something) ≈ 8-10. The values are moderate. Simple implementation might be fine.
However, I need to match PyTorch's output. PyTorch uses a numerically stable implementation. Let me just match it closely.
Actually, let me check: LSE of clamped values in [-10, 10] with 8192 elements. LSE ≈ max + log(8192) ≈ 10 + 9 = 19 if all values are 10. Or if values are random after matmul... well, inputs are rand(1024, 8192), so initial values are 0-1. After matmul with random weights... the output values are sums of 8192 uniform[0,1] * random weights. This could have a wide range. But after scaling by 4, they could be large. Then clamp to [-10, 10].
If many values are near 10, LSE ≈ 10 + log(8192) ≈ 19.0.
For mish of 19: softplus(19) ≈ 19, tanh(19) ≈ 1.0, mish(19) ≈ 19 * 1.0 = 19. So result = 19 * 19 = 361.
For numerical stability with moderate values like 19, a simple implementation should be fine. But to be safe and match PyTorch:
```cpp
__device__ inline float mish(float x) {
float sp = x > 20.0f ? x : log1pf(expf(x));
return x * tanhf(sp);
}
```
Actually, PyTorch's softplus implementation uses `x > 0 ? x + log1p(exp(-x)) : log1p(exp(x))`. Let me use that.
```cpp
__device__ inline float softplus(float x) {
return x > 0.0f ? x + log1pf(expf(-x)) : log1pf(expf(x));
}
__device__ inline float mish(float x) {
return x * tanhf(softplus(x));
}
```
Wait, for x = 19: exp(-19) is tiny, log1p(exp(-19)) ≈ exp(-19), so softplus(19) ≈ 19 + exp(-19) ≈ 19. This is fine.
For x = -19: exp(-19) is tiny, softplus(-19) = log1p(exp(-19)) ≈ exp(-19) ≈ 0. This is fine.
OK, I'll use a stable softplus.
Now let me write the complete code:
Actually, one more thought: can I fuse the mish computation into the LSE kernel? The LSE kernel outputs 1 value per row. The mish kernel operates on those values. Since there are only 1024 values, the mish kernel is tiny. I'd rather keep it separate for simplicity, but I could also have the LSE kernel do the mish. Since each block computes LSE for one row, after computing LSE, thread 0 could compute mish and write the final result. This saves launching another kernel.
Let me do that!
```cpp
__global__ void fused_scale_clamp_lse_mish_kernel(...) {
// ... compute LSE ...
if (threadIdx.x == 0) {
float lse = row_max + logf(row_sum);
float sp = lse > 0.0f ? lse + log1pf(expf(-lse)) : log1pf(expf(lse));
float m = lse * tanhf(sp);
output[row] = m;
}
}
```
This fuses everything into one kernel!
Now the full forward pass is:
1. `y = F.linear(x, self.weight, self.bias)` # (1024, 8192)
2. `output = fused_scale_clamp_lse_mish_kernel(y)` # (1024, 1)
This is very clean.
Wait, I need to verify: the reference does `x = x * torch.nn.functional.mish(x)`. That is, the LSE result multiplied by mish(LSE result). So it's `lse * mish(lse)`.
In my notation: `m = lse * tanh(softplus(lse))` which is `lse * mish(lse)`. Yes, that's correct.
Now let me write the complete CUDA code.
One more optimization: I should check if using `block_reduce_max` and `block_reduce_sum` with `__syncthreads()` in between is correct. Yes.
For the shared memory approach, each block needs `hidden_size * sizeof(float)` bytes. With batch_size=1024 blocks, each using up to 32KB shared memory, that's fine.
Let me also consider whether to use static or dynamic shared memory. Dynamic is easier for varying sizes.
Complete kernel:
```cpp
#include <torch/extension.h>
#include <cuda_runtime.h>
#define BLOCK_SIZE 1024
__inline__ __device__ float warp_reduce_max(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
return val;
}
__inline__ __device__ float warp_reduce_sum(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
return val;
}
__inline__ __device__ float block_reduce_max(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_max(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -INFINITY;
if (wid == 0) val = warp_reduce_max(val);
return val;
}
__inline__ __device__ float block_reduce_sum(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
if (wid == 0) val = warp_reduce_sum(val);
return val;
}
__device__ inline float stable_softplus(float x) {
return x > 0.0f ? x + log1pf(expf(-x)) : log1pf(expf(x));
}
__global__ void fused_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int hidden_size,
float scale,
float clamp_min,
float clamp_max
) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_ptr = input + row * hidden_size;
int tid = threadIdx.x;
extern __shared__ float sdata[];
int elems_per_thread = (hidden_size + BLOCK_SIZE - 1) / BLOCK_SIZE;
int start = tid * elems_per_thread;
int end = start + elems_per_thread;
if (end > hidden_size) end = hidden_size;
// Phase 1: load, scale, clamp, store to shared, find local max
float local_max = -INFINITY;
for (int idx = start; idx < end; idx++) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
sdata[idx] = val;
local_max = fmaxf(local_max, val);
}
__syncthreads();
float row_max = block_reduce_max(local_max);
__syncthreads();
// Phase 2: compute sum of exp(x - max)
float local_sum = 0.0f;
for (int idx = start; idx < end; idx++) {
local_sum += expf(sdata[idx] - row_max);
}
float row_sum = block_reduce_sum(local_sum);
if (tid == 0) {
float lse = row_max + logf(row_sum);
float sp = stable_softplus(lse);
float mish = lse * tanhf(sp);
output[row] = mish;
}
}
torch::Tensor forward(torch::Tensor input, float scale, float clamp_min, float clamp_max) {
int batch_size = input.size(0);
int hidden_size = input.size(1);
auto output = torch::empty({batch_size, 1}, input.options());
dim3 grid(batch_size);
dim3 block(BLOCK_SIZE);
size_t shared_mem_size = hidden_size * sizeof(float);
fused_kernel<<<grid, block, shared_mem_size>>>(
input.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
hidden_size,
scale,
clamp_min,
clamp_max
);
return output;
}
```
Wait, there's a subtle issue with `block_reduce_max`. It uses `__syncthreads()` inside. If I call it, then `__syncthreads()` after it, that's fine. But what about the phase 2 where I call `block_reduce_sum`? It also has `__syncthreads()` inside. And there are `__syncthreads()` between phases. This should be fine.
Actually, there is a problem! In `block_reduce_max`, the first `__syncthreads()` is reached by all threads. But then:
```cpp
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -INFINITY;
if (wid == 0) val = warp_reduce_max(val);
```
For threads where `wid != 0`, they don't participate in the second warp reduce, but `warp_reduce_max` uses `__shfl_down_sync(0xFFFFFFFF, ...)` which includes all threads in the warp. Actually, for threads where `wid != 0`, they are still in a warp with other threads. The shuffle operations happen within the warp. For threads with `wid != 0` and `lane == 0`... wait, let me reconsider.
`threadIdx.x` ranges from 0 to 1023.
- `wid = threadIdx.x >> 5` ranges from 0 to 31.
- `lane = threadIdx.x & 31` ranges from 0 to 31.
For `threadIdx.x = 32`, `wid = 1`, `lane = 0`.
In the second part of `block_reduce_max`:
```cpp
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -INFINITY;
```
`blockDim.x >> 5` = 1024 >> 5 = 32. So `threadIdx.x < 32` is true only for threads 0-31. For thread 32, `val = -INFINITY`.
Then `if (wid == 0) val = warp_reduce_max(val);`
For thread 32, `wid = 1`, so it doesn't enter the if. `val` remains -INFINITY.
But wait, `warp_reduce_max` uses `__shfl_down_sync(0xFFFFFFFF, val, offset)`. This is a warp-level primitive. For threads in warp 1 (tid 32-63), they all have `wid = 1`, so they all skip the `if (wid == 0)` block. Their `val` remains -INFINITY. But since they're all in the same warp and don't call `__shfl_down_sync`, is there a problem?
Actually, in CUDA, `__shfl_down_sync` is only called by threads with `wid == 0`. Threads with `wid != 0` don't call it. Is this a problem? No, because the mask `0xFFFFFFFF` is for the warp that the thread is in. But wait, there's a subtlety: threads with `wid != 0` don't participate in the shuffle, but they also don't need the result. For threads in warp 0, they all call `warp_reduce_max` and get the same result. For threads in other warps, they don't call it.
Actually, this is fine because:
1. Threads with `wid != 0` don't use `val` after this point in `block_reduce_max` (they just return it, but the caller only uses the result from thread 0 typically).
2. The shuffle happens within each warp independently. Warp 0 does its shuffle, other warps don't.
Wait, but CUDA requires all threads in a warp to execute the same instruction (or be masked out). If some threads in a warp hit the `if (wid == 0)` and others don't... but all threads in warp 0 have `wid == 0`, so they all enter. All threads in warp 1 have `wid == 1`, so they all don't enter. Since the branch is uniform within each warp, there's no divergence issue.
So `block_reduce_max` is fine.
Similarly for `block_reduce_sum`.
One more concern: the first `__syncthreads()` in `block_reduce_max` is inside the function. After I call `block_reduce_max`, I have `__syncthreads()` again in the kernel. That's fine - double syncthreads is allowed as long as all threads reach both.
Actually, there's a potential issue: after `block_reduce_max`, only thread 0 (or threads in warp 0, lane 0) has the correct result. Actually no - `block_reduce_max` uses warp reduce, and the second warp reduce writes the result to `val` for threads in warp 0. Specifically, thread 0 gets the final result.
But in the kernel, I don't need `row_max` in all threads for phase 2. I use it in phase 2 where each thread computes `expf(sdata[idx] - row_max)`. For this, all threads need `row_max`.
Looking at `block_reduce_max`: after the function, which threads have the correct `row_max`?
- For warp 0 threads (tid 0-31), after `warp_reduce_max(val)`, all threads in warp 0 have the final max value in `val` (because warp broadcast).
- For other warps, they don't.
So in the kernel, after `row_max = block_reduce_max(local_max);`, threads in warp 0 have the correct value, but threads in other warps don't!
I need to broadcast `row_max` to all threads. I can use shared memory for this:
```cpp
float row_max = block_reduce_max(local_max);
__shared__ float s_row_max;
if (threadIdx.x == 0) {
s_row_max = row_max;
}
__syncthreads();
row_max = s_row_max;
```
Wait, but `block_reduce_max` already has a `__syncthreads()`. And I need another one to ensure `s_row_max` is written before read.
Let me modify `block_reduce_max` to NOT do `__syncthreads()` at the end, or add the broadcast in the kernel.
Actually, let me just handle it in the kernel:
```cpp
float row_max = block_reduce_max(local_max);
__shared__ float s_row_max;
if (threadIdx.x == 0) {
s_row_max = row_max;
}
__syncthreads();
row_max = s_row_max;
```
Wait, `block_reduce_max` has its own `__syncthreads()`. After that function returns, all threads are synced. Then thread 0 writes to `s_row_max`. Then I need another `__syncthreads()` for other threads to read it. Then `row_max = s_row_max`. This is correct.
Similarly for `row_sum` - but `row_sum` is only needed by thread 0 at the end, so no broadcast needed.
Actually, wait - `block_reduce_sum` also has `__syncthreads()` inside. After it returns, thread 0 has the value. That's fine because only thread 0 uses it.
Let me rewrite the reduction functions to not include syncthreads if it makes it cleaner... Actually, the current approach is fine.
But wait, looking more carefully at `block_reduce_max` and `block_reduce_sum`, they both write to `shared[wid]`. At different times, so the shared memory is reused. But both use `__syncthreads()`. Let me make sure `shared` is declared inside each function - yes, it is. Each function has its own `__shared__ float shared[32]`. These should be in different memory locations or the compiler handles them. Actually, static shared memory is just one segment, and the arrays are overlaid if not used simultaneously. Since they're in different functions and used at different times, this should be fine with sufficient syncthreads between them.
Actually, CUDA might place both arrays at the same offset in shared memory since they're in different scopes and not live simultaneously. This is fine.
Let me now put it all together in a Python module.
I want to verify: in the reference, `self.scale_factor = 2.0`, and `x = x * scale_factor` then `x = x + x`. So total scaling is 2.0 * 2.0 = 4.0. I'll pass 4.0 to the kernel.
Let me write the ModelNew class:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
#define BLOCK_SIZE 1024
__inline__ __device__ float warp_reduce_max(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
return val;
}
__inline__ __device__ float warp_reduce_sum(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
return val;
}
__inline__ __device__ float block_reduce_max(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_max(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -INFINITY;
if (wid == 0) val = warp_reduce_max(val);
return val;
}
__inline__ __device__ float block_reduce_sum(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
if (wid == 0) val = warp_reduce_sum(val);
return val;
}
__device__ inline float stable_softplus(float x) {
return x > 0.0f ? x + log1pf(expf(-x)) : log1pf(expf(x));
}
__global__ void fused_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int hidden_size,
float scale,
float clamp_min,
float clamp_max
) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_ptr = input + row * hidden_size;
int tid = threadIdx.x;
extern __shared__ float sdata[];
int elems_per_thread = (hidden_size + BLOCK_SIZE - 1) / BLOCK_SIZE;
int start = tid * elems_per_thread;
int end = start + elems_per_thread;
if (end > hidden_size) end = hidden_size;
// Phase 1: load, scale, clamp, store to shared, find local max
float local_max = -INFINITY;
for (int idx = start; idx < end; idx++) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
sdata[idx] = val;
local_max = fmaxf(local_max, val);
}
__syncthreads();
float row_max = block_reduce_max(local_max);
__shared__ float s_row_max;
if (tid == 0) s_row_max = row_max;
__syncthreads();
row_max = s_row_max;
// Phase 2: compute sum of exp(x - max)
float local_sum = 0.0f;
for (int idx = start; idx < end; idx++) {
local_sum += expf(sdata[idx] - row_max);
}
float row_sum = block_reduce_sum(local_sum);
if (tid == 0) {
float lse = row_max + logf(row_sum);
float sp = stable_softplus(lse);
float mish_val = lse * tanhf(sp);
output[row] = mish_val;
}
}
torch::Tensor forward(torch::Tensor input, float scale, float clamp_min, float clamp_max) {
int batch_size = input.size(0);
int hidden_size = input.size(1);
auto output = torch::empty({batch_size, 1}, input.options());
dim3 grid(batch_size);
dim3 block(BLOCK_SIZE);
size_t shared_mem_size = hidden_size * sizeof(float);
fused_kernel<<<grid, block, shared_mem_size>>>(
input.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
hidden_size,
scale,
clamp_min,
clamp_max
);
return output;
}
'''
cpp_source = '''
torch::Tensor forward(torch::Tensor input, float scale, float clamp_min, float clamp_max);
'''
# Compile
fused_op = load_inline(
name="fused_op",
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=["forward"],
extra_cuda_cflags=["-O3"],
)
class ModelNew(nn.Module):
def __init__(self, input_size, hidden_size, scale_factor, clamp_min, clamp_max):
super(ModelNew, self).__init__()
self.hidden_size = hidden_size
self.scale_factor = scale_factor * 2.0 # account for x = x + x
self.clamp_min = clamp_min
self.clamp_max = clamp_max
self.weight = nn.Parameter(torch.empty(hidden_size, input_size))
self.bias = nn.Parameter(torch.empty(hidden_size))
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
nn.init.uniform_(self.bias, -bound, bound)
def forward(self, x):
y = F.linear(x, self.weight, self.bias)
return fused_op.forward(y, self.scale_factor, self.clamp_min, self.clamp_max)
```
Wait, `self.scale_factor = scale_factor * 2.0` because `x = x * scale_factor` (2.0) and `x = x + x` (another 2x). So total is 4.0.
For the initialization, I need to match PyTorch's Linear initialization. nn.Linear uses `kaiming_uniform_` with `a=math.sqrt(5)` for weights, and bias init as I wrote.
Actually, looking at PyTorch source for Linear:
```python
def reset_parameters(self):
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
init.uniform_(self.bias, -bound, bound)
```
Yes, that's what I have. But wait, in `nn.Linear`, the weight shape is `(out_features, in_features) = (hidden_size, input_size)`. And `F.linear(x, weight, bias)` expects weight of shape `(out_features, in_features)`. So my weight shape `torch.empty(hidden_size, input_size)` is correct.
I also need to import math.
Let me also check: when I use `load_inline`, the first argument is `name` which should be unique. I'll use something like `fused_custom_kernel`.
One issue: `load_inline` may fail if there's a syntax error or CUDA compilation error. Let me make sure everything is correct.
Actually, I realize there's a potential issue with `extern __shared__ float sdata[];`. In CUDA, when using `extern __shared__`, the array must be declared at the kernel level or as a parameter (for dynamic shared memory). What I have is correct.
But wait, for static shared memory, if I declare `__shared__ float sdata[8192];` inside the kernel, it's fine. But 8192 is not a compile-time constant for a general kernel... though for this problem it's fixed. However, using dynamic shared memory is more flexible.
Let me also consider whether BLOCK_SIZE=1024 is optimal. With 1024 threads and 8192 hidden_size, each thread handles 8 elements. This is fine. Could 512 threads with 16 elements each be better? Maybe slightly fewer registers per block, allowing more blocks per SM. But 1024 is a standard choice.
Actually, for H100, we might want to consider using 512 threads to increase occupancy. With 1024 threads and some registers, we might only get 1 block per SM. With 512 threads, we might get 2 or more.
Let me estimate register usage. Each thread has some local variables: `local_max`, `local_sum`, `row_max`, `row_sum`, `start`, `end`, `idx`, `val`, plus loop counters. Maybe 20-30 registers. At 1024 threads * 30 registers = 30K. H100 has 64K, so we can fit 2 blocks. With shared memory at 32KB per block, 2 blocks = 64KB, still under 228KB. So we get 2 blocks per SM, which gives 2048 threads, maxing out the SM thread capacity. So occupancy is already good.
Wait, actually H100 has 64K 32-bit registers per SM. 2 blocks of 1024 threads with 32 registers each = 64K registers. That exactly fills the register file. But if each thread uses more than 32 registers, we might only fit 1 block. Let me check with the compiler.
To be safe, I can reduce BLOCK_SIZE to 512. Then even with 64 registers per thread, we'd have 512 * 64 = 32K, so we can fit 2 blocks. And 2 * 512 = 1024 threads, which is half capacity. Hmm, that might be worse.
With BLOCK_SIZE=1024 and 30 registers: 30K registers per block. Can fit 2 blocks = 60K < 64K, giving 2048 threads. That's actually max threads per SM. So 100% occupancy (in terms of threads).
But if it's 33 registers: 33K per block. 2 blocks = 66K > 64K. So only 1 block = 1024 threads = 50% occupancy.
With 512 threads and 33 registers: 16.9K per block. 3 blocks = 50.6K < 64K. 3 * 512 = 1536 threads = 75% occupancy. But 1536 > 1024 in terms of useful warps.
Actually, the number of warps matters more than thread count for latency hiding. With 1024 threads at 50% occupancy = 32 warps. With 512 threads at 75% occupancy but wait, 3 blocks of 512 = 1536 threads = 48 warps. 48 warps is better than 32 warps for occupancy.
But the kernel is mostly memory-bound. More warps help hide latency. So 512 threads might be better if it allows more warps per SM.
However, with 8192 hidden_size and 512 threads, each thread handles 16 elements. The loops are longer but instruction count per thread is similar.
I think 1024 is fine for now. If occupancy is low, I can adjust later.
Let me also check: the matmul is the dominant operation anyway. The fused kernel is secondary. So as long as it's reasonably efficient, the main win comes from fusing the operations and reducing memory traffic.
Actually, wait. The reference does:
1. matmul -> writes (1024, 8192) = 32MB
2. scale -> reads 32MB, writes 32MB
3. add -> reads 32MB, writes 32MB
4. clamp -> reads 32MB, writes 32MB
5. logsumexp -> reads 32MB, writes 4KB
6. mish multiplication -> reads 4KB, writes 4KB
Total memory traffic: matmul writes 32MB + ~160MB for the element-wise ops = ~192MB.
My approach:
1. matmul -> writes 32MB (but internally cuBLAS might use more, but that's the same for both)
2. fused kernel -> reads 32MB, writes 4KB
So I save ~160MB of memory traffic! That's a huge win.
For time: at 3.35 TB/s, 160MB takes ~0.048 ms. Matmul takes ~0.14 ms. Total reference might be ~0.19 ms + kernel launch overhead. My version might be ~0.15 ms. That's about 1.27x speedup.
But I should also check if PyTorch fuses some of these operations. PyTorch has a JIT compiler that might fuse some ops, but logsumexp and mish are not typically fused. PyTorch might fuse scale + add + clamp into one kernel, but not with logsumexp.
So my approach should be significantly faster.
One more thing: I should make sure `forward` returns `output` with shape `(batch_size, 1)`. The reference returns this shape.
Let me finalize my code and compile it.
Actually, one more thought. Can I also fuse the bias addition and scaling into the matmul using torch.addmm? `torch.addmm(bias, x, weight.t())` or `F.linear` which calls addmm internally. This is already optimal.
But wait, I can use `torch.matmul` and then `torch.add` - no, `F.linear` is the standard way and should call optimized cuBLAS.
Actually, let me use `torch.nn.functional.linear` which is the same as the reference.
Wait, the reference uses `self.matmul = nn.Linear(...)`. In its forward, it calls `self.matmul(x)`. This calls the Linear layer which does the matmul and bias add.
In my ModelNew, I can either:
1. Create an `nn.Linear` layer and call it
2. Manually do `F.linear(x, self.weight, self.bias)`
Option 2 is fine and matches. But if I create the parameters manually, I need to make sure they're named correctly and initialized the same way.
Actually, I can just do:
```python
self.matmul = nn.Linear(input_size, hidden_size)
```
in my ModelNew, and then call it. This is the simplest and ensures exact same initialization and behavior.
```python
class ModelNew(nn.Module):
def __init__(self, input_size, hidden_size, scale_factor, clamp_min, clamp_max):
super(ModelNew, self).__init__()
self.matmul = nn.Linear(input_size, hidden_size)
self.scale_factor = scale_factor * 2.0
self.clamp_min = clamp_min
self.clamp_max = clamp_max
def forward(self, x):
y = self.matmul(x)
return fused_op.forward(y, self.scale_factor, self.clamp_min, self.clamp_max)
```
This is much simpler and avoids initialization mismatches.
Wait, but if I use `nn.Linear`, the weight will be transposed compared to what `F.linear` expects? No, `nn.Linear` stores weight as `(out_features, in_features)`, and its forward calls `F.linear(input, self.weight, self.bias)`. So it's exactly the same.
Let me use this simpler approach.
Final code:
```python
import torch
import torch.nn as nn
import math
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
#define BLOCK_SIZE 1024
__inline__ __device__ float warp_reduce_max(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
return val;
}
__inline__ __device__ float warp_reduce_sum(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
return val;
}
__inline__ __device__ float block_reduce_max(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_max(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -INFINITY;
if (wid == 0) val = warp_reduce_max(val);
return val;
}
__inline__ __device__ float block_reduce_sum(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
if (wid == 0) val = warp_reduce_sum(val);
return val;
}
__device__ inline float stable_softplus(float x) {
return x > 0.0f ? x + log1pf(expf(-x)) : log1pf(expf(x));
}
__global__ void fused_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int hidden_size,
float scale,
float clamp_min,
float clamp_max
) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_ptr = input + row * hidden_size;
int tid = threadIdx.x;
extern __shared__ float sdata[];
int elems_per_thread = (hidden_size + BLOCK_SIZE - 1) / BLOCK_SIZE;
int start = tid * elems_per_thread;
int end = start + elems_per_thread;
if (end > hidden_size) end = hidden_size;
// Phase 1: load, scale, clamp, store to shared, find local max
float local_max = -INFINITY;
for (int idx = start; idx < end; idx++) {
float val = row_ptr[idx] * scale;
val = fminf(fmaxf(val, clamp_min), clamp_max);
sdata[idx] = val;
local_max = fmaxf(local_max, val);
}
__syncthreads();
float row_max = block_reduce_max(local_max);
__shared__ float s_row_max;
if (tid == 0) s_row_max = row_max;
__syncthreads();
row_max = s_row_max;
// Phase 2: compute sum of exp(x - max)
float local_sum = 0.0f;
for (int idx = start; idx < end; idx++) {
local_sum += expf(sdata[idx] - row_max);
}
float row_sum = block_reduce_sum(local_sum);
if (tid == 0) {
float lse = row_max + logf(row_sum);
float sp = stable_softplus(lse);
float mish_val = lse * tanhf(sp);
output[row] = mish_val;
}
}
torch::Tensor forward(torch::Tensor input, float scale, float clamp_min, float clamp_max) {
int batch_size = input.size(0);
int hidden_size = input.size(1);
auto output = torch::empty({batch_size, 1}, input.options());
dim3 grid(batch_size);
dim3 block(BLOCK_SIZE);
size_t shared_mem_size = hidden_size * sizeof(float);
fused_kernel<<<grid, block, shared_mem_size>>>(
input.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
hidden_size,
scale,
clamp_min,
clamp_max
);
return output;
}
'''
cpp_source = '''
torch::Tensor forward(torch::Tensor input, float scale, float clamp_min, float clamp_max);
'''
fused_op = load_inline(
name="fused_custom_kernel",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["forward"],
extra_cuda_cflags=["-O3"],
extra_cflags=["-O3"],
)
class ModelNew(nn.Module):
def __init__(self, input_size, hidden_size, scale_factor, clamp_min, clamp_max):
super(ModelNew, self).__init__()
self.matmul = nn.Linear(input_size, hidden_size)
self.scale_factor = scale_factor * 2.0
self.clamp_min = clamp_min
self.clamp_max = clamp_max
def forward(self, x):
y = self.matmul(x)
return fused_op.forward(y, self.scale_factor, self.clamp_min, self.clamp_max)
```
Wait, I need to make sure I didn't miss anything. Let me re-check `load_inline` signature.
`load_inline(name, cpp_sources, cuda_sources, functions, extra_cuda_cflags, extra_cflags, ...)`
Actually, I'm not sure about the parameter names. Let me check standard usage:
```python
fused_op = load_inline(
name="fused_custom_kernel",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["forward"],
extra_cuda_cflags=["-O3"],
extra_cflags=["-O3"],
)
```
`cpp_sources` can be a string or list of strings. Same for `cuda_sources`. `functions` is a list of function names. This looks correct.
One potential issue: the `cpp_source` declaration might not be needed if we're only exporting the CUDA function. Actually, `load_inline` generates a pybind11 module. The `cpp_sources` provides any additional C++ code, and `cuda_sources` provides the CUDA code. The function `forward` is a host function that calls the kernel, so it needs to be in the compiled code. Since I defined it in the `cuda_sources`, it should be available. The `cpp_sources` in this case is just the forward declaration, which might not be necessary.
Hmm, actually in `load_inline`, functions declared in `cuda_sources` with `torch::Tensor` return types are compiled with nvcc and should be exported. Let me check if I need `cpp_sources` at all.
Looking at examples of `load_inline`:
```python
op = load_inline(
name="my_op",
cpp_sources="torch::Tensor my_op(torch::Tensor x);",
cuda_sources=cuda_source,
functions=["my_op"],
)
```
The `cpp_sources` typically contains just the declaration. The implementation is in `cuda_sources`. This is because nvcc compiles the CUDA sources and links them with the C++ wrapper.
Actually, I've seen both ways. To be safe, I'll keep the declaration in `cpp_sources` and implementation in `cuda_sources`.
Now, let me compile and test this.
Actually, I realize there might be an issue with the `cpp_sources` parameter name. Is it `cpp_sources` or `cpp_source`? Let me check.
Looking at PyTorch docs: `torch.utils.cpp_extension.load_inline(name, cpp_sources, cuda_sources, functions, extra_cuda_cflags=None, extra_cflags=None, verbose=False, with_cuda=None, is_python_module=True, keep_intermediates=True)`
Yes, it's `cpp_sources` and `cuda_sources`. Both can be strings or lists of strings.
One issue: my CUDA code uses `log1pf` and `expf`. These should be available in CUDA math library.
Another thing: I should include `-std=c++17` in cflags. And for H100, maybe `-arch=sm_90`.
Actually, `load_inline` typically auto-detects the architecture. But I can specify it explicitly.
Let me add `extra_cuda_cflags=["-O3", "--use_fast_math"]`. `--use_fast_math` might speed up expf, logf, tanhf at the cost of some precision. But the tolerance is 1e-4 for fp32, so fast math should be fine.
Actually, let me be conservative and not use `--use_fast_math` initially. If it's too slow, I can add it later.
Let me compile now.
Wait, I need to be careful about one thing: the CUDA code string in Python should be properly formatted. Since it's a triple-quoted string, indentation matters for readability but not for the C++ compiler. However, I need to make sure there are no issues with backslashes or special characters.
Let me write the complete code.