reasoning
This is a simple Softplus activation on a large tensor: [4096, 393216] = 1,610,612,736 elements, about 6.4 GB of fp32 data (input + output).
The operation is elementwise: softplus(x) = (1/beta) * log(1 + exp(beta * x)). Default beta=1, threshold=20 in PyTorch.
For x*beta > threshold, PyTorch returns x directly (to avoid overflow in exp). Default threshold is 20, beta is 1.
This is a memory-bound operation. We need to maximize memory bandwidth with a simple fused CUDA kernel. The H100 has 3.35 TB/s bandwidth. For 6.4 GB, theoretical minimum is about 1.9 ms. We should aim for close to that.
We should write a kernel that loads float, computes softplus efficiently, and stores float.
Given the large size (1.6B elements), we can use a straightforward grid-stride loop with each thread handling multiple elements.
We need to use `load_inline` with a C++ wrapper. Let's write a simple kernel.
One thing to be careful about: we need to define a custom CUDA kernel in `cuda_sources` and a host wrapper in `cpp_sources`.
Since it's elementwise, we should also use `__restrict__` for pointer arguments and perhaps `__builtin_assume` for alignment if possible. But standard alignment should be fine.
The kernel:
```cpp
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void softplus_kernel(const float* __restrict__ input, float* __restrict__ output, int64_t n) {
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
const int64_t stride = gridDim.x * blockDim.x;
for (; idx < n; idx += stride) {
float x = input[idx];
// softplus(x) = log(1 + exp(x)) for x * beta <= threshold, else x
// beta=1, threshold=20
float y;
if (x > 20.0f) {
y = x;
} else {
y = log1pf(expf(x)); // log1p(exp(x)) is more stable
}
output[idx] = y;
}
}
torch::Tensor softplus_cuda(torch::Tensor input) {
auto output = torch::empty_like(input);
int64_t n = input.numel();
const int threads = 512;
const int blocks = std::min((n + threads - 1) / threads, (int64_t)65535 * 32); // large enough
softplus_kernel<<<blocks, threads>>>(input.data_ptr<float>(), output.data_ptr<float>(), n);
return output;
}
```
Wait, we should use `at::cuda::CUDAStreamGuard` or just rely on the default stream. `torch::empty_like` should be fine.
For performance, we should consider vectorized loads/stores. Since H100 has good L2 cache and large bandwidth, vectorized memory access (float4) could help. However, with such a large tensor and purely memory-bound, the memory controller should handle coalescing well with simple float loads. But vectorized loads can reduce instruction overhead.
Let's try a simple vectorized version using float4.
A float4 load/store handles 4 elements per thread, reducing the number of instructions and potentially improving bandwidth.
```cpp
__global__ void softplus_kernel_float4(const float* __restrict__ input, float* __restrict__ output, int64_t n) {
int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
const int64_t stride = gridDim.x * blockDim.x * 4;
for (; idx < n; idx += stride) {
float4 in = reinterpret_cast<const float4*>(input)[idx / 4];
float4 out;
// Process each element
// softplus logic
// ...
reinterpret_cast<float4*>(output)[idx / 4] = out;
}
}
```
But vectorized loads require alignment. PyTorch tensors for CUDA are generally aligned to 16 bytes or more for large sizes, but it's not guaranteed for all dimensions. However, a contiguous 1D tensor of size `batch_size * dim` of float is likely aligned to at least 4 bytes. For float4 alignment, we need 16 bytes. PyTorch's `torch.empty` usually aligns to 256 bytes, so it should be safe for the start of the tensor. But if the tensor is sliced or not contiguous, alignment might be an issue. Since the reference implementation just calls `torch.nn.functional.softplus(x)`, and `x` comes from `torch.rand(batch_size, dim)`, it is contiguous and aligned. Using float4 should be safe, but we should probably add a check or fallback. Actually, it's simpler to just do scalar loads. The compiler might vectorize anyway with `-O3`. Let's start with a simple scalar kernel and see the performance.
Wait, the grid size should be large enough to saturate the GPU. For H100 with 132 SMs (actually H100 SXM has 132 SMs), we need many blocks. With 512 threads per block, 65535 blocks * 512 = 33M threads, which is way more than enough. Actually, we should compute blocks as something like `(n + threads - 1) / threads`. For 1.6B elements, that's about 3.1M blocks. That might be too many blocks for a 1D grid (max is 2^31-1, so it's fine), but each block does very little work. The grid-stride loop handles that. We should set `blocks` to a reasonable number, e.g., the number of SMs * some factor (e.g., 8 or 16), and let the grid-stride loop handle the rest. That's usually better for elementwise kernels to reduce launch overhead.
Number of SMs for H100: 132 (SXM) or 114 (PCIe). The spec just says H100 80GB HBM3. H100 80GB HBM3 is usually SXM5 with 132 SMs. Let's assume 132. A typical good block count is something like 132 * 16 = 2112 blocks. Or we can just launch `(n + threads - 1) / threads` and rely on the CUDA runtime. For very large grids, launching millions of blocks is fine on modern GPUs. However, having too many blocks can sometimes slightly increase overhead. Let's use a fixed large number of blocks like `(n + threads - 1) / threads` capped at maybe 2^20 or similar.
Actually, many optimized elementwise kernels in PyTorch use `num_threads` and `num_blocks` computed as `(N + num_threads - 1) / num_threads`. For 1.6B elements and 512 threads, that's ~3.2M blocks. Modern GPUs can handle that. I think it's fine.
But let's also check occupancy. With 512 threads, registers might limit occupancy but we won't use many registers.
Let me write the kernel code.
Important: `load_inline` requires `cuda_sources` string and `cpp_sources` string. We must not include `PYBIND11_MODULE`. We must use `functions=['softplus_cuda']`.
Let me draft the full Python module.
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void softplus_kernel(const float* __restrict__ input, float* __restrict__ output, int64_t n) {
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
const int64_t stride = gridDim.x * blockDim.x;
for (; idx < n; idx += stride) {
float x = input[idx];
float y;
if (x > 20.0f) {
y = x;
} else {
y = log1pf(expf(x));
}
output[idx] = y;
}
}
torch::Tensor softplus_cuda(torch::Tensor input) {
auto output = torch::empty_like(input);
int64_t n = input.numel();
const int threads = 512;
const int blocks = std::min((n + threads - 1) / threads, (int64_t)65535 * 32);
softplus_kernel<<<blocks, threads>>>(input.data_ptr<float>(), output.data_ptr<float>(), n);
return output;
}
'''
cpp_source = '''
torch::Tensor softplus_cuda(torch::Tensor input);
'''
module = load_inline(
name='softplus_ext',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['softplus_cuda'],
extra_cuda_cflags=['-O3', '--use_fast_math'],
verbose=False,
)
class ModelNew(torch.nn.Module):
def __init__(self):
super(ModelNew, self).__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return module.softplus_cuda(x)
```
Wait, we don't want `torch.extension.h` in the CUDA source? It's fine, it's needed for `torch::Tensor`. But `cuda_sources` doesn't necessarily need `torch/extension.h` if the kernel itself doesn't use torch types. However, the wrapper function uses `torch::Tensor` which is in `torch/extension.h`. In `load_inline`, the host wrapper is in the C++ source usually, but we can put it in `cuda_sources` (compiled with nvcc) since nvcc can compile C++ code containing `torch::Tensor`. That's the common approach: put everything in `cuda_sources`. Actually, `cpp_sources` is compiled with the C++ compiler and `cuda_sources` with nvcc. Since the function uses CUDA kernel launches, it must be in a `.cu` file (i.e., `cuda_sources`). So we can put the wrapper in `cuda_sources` and leave `cpp_sources` with just the declaration.
But `load_inline` needs the declaration in `cpp_sources` to generate the pybind11 wrapper. Actually, I think `functions=[...]` tells it to look for those functions in the compiled sources. It generates a pybind11 wrapper that calls them. The `cpp_sources` can be empty or contain the forward declarations. Typically, you put the declarations in `cpp_sources` and definitions in `cuda_sources`.
Now, `extra_cuda_cflags` with `--use_fast_math` might affect numerical precision. `log1pf` and `expf` with `--use_fast_math` have slight differences. The tolerance for fp32 is 1e-4 (atol and rtol). `--use_fast_math` usually changes results by up to a few ULPs, but let's check if it's within tolerance. To be safe, maybe don't use `--use_fast_math` initially. But for performance, `--use_fast_math` might allow faster `expf` and `log1pf`. Let's first try without it to ensure correctness.
Also, we should handle the beta and threshold parameters? The reference `torch.nn.functional.softplus(x)` uses default `beta=1`, `threshold=20`. The problem says "replace the forward computation", and the model has no parameters. So we just need to match the default.
Let's think about numerical stability. `log1p(exp(x))` is the standard way. For x > 20, PyTorch returns x. So we implement that.
Actually, PyTorch's softplus is:
`softplus(x) = (1/beta) * log(1 + exp(beta * x))` if `beta * x <= threshold`, else `x`.
With default beta=1, threshold=20.
So our implementation is correct.
One thing: PyTorch may use `beta` and `threshold` as parameters. Since the `Model` does not accept them and uses defaults, we are fine.
Another idea: instead of writing a custom kernel, could we use `torch.nn.functional.softplus` with `torch.jit.script` or something? No, we need a custom CUDA kernel.
Another optimization: use vectorized memory access (float4) to saturate bandwidth better. Let's try to implement a float4 version with alignment check and scalar fallback for tail.
But wait, the simplest kernel might already be very close to the reference. The reference is just calling `torch.nn.functional.softplus`, which is an ATen op implemented on CUDA. That op is likely already quite well optimized. However, `softplus` might not be as optimized as a simple elementwise kernel; it might use the standard `gpuUnaryKernel` machinery which is good but might have some overhead. Our custom kernel might be slightly faster if we tune block size and use fast math. But it's possible the reference is already near-optimal.
Let's see if we can do better by fusing or vectorizing. Since the operation is elementwise, the best we can do is saturate memory bandwidth.
Let's write a vectorized kernel using float4 and see.
Vectorized kernel:
```cpp
__global__ void softplus_kernel_vec4(const float* __restrict__ input, float* __restrict__ output, int64_t n) {
int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
const int64_t stride = gridDim.x * blockDim.x * 4;
for (; idx < n; idx += stride) {
float4 in = reinterpret_cast<const float4*>(input)[idx / 4];
float4 out;
#pragma unroll
for (int i = 0; i < 4; ++i) {
float x = (&in.x)[i];
float y;
if (x > 20.0f) {
y = x;
} else {
y = log1pf(expf(x));
}
(&out.x)[i] = y;
}
reinterpret_cast<float4*>(output)[idx / 4] = out;
}
}
```
Then we need a tail scalar loop for remaining elements. We could launch the vectorized kernel with `n4 = n / 4` and then a scalar kernel for the remainder. But we can also just handle the remainder in the same kernel by checking bounds.
Simpler: the main loop processes `n / 4` elements using float4, and the tail is handled by a separate scalar kernel or by the same threads but with bounds checks. Actually, if we have a grid-stride loop with `stride` such that we cover all `n` elements, the last iteration might go out of bounds. So we need to check `idx + 3 < n` for the float4 load.
```cpp
for (; idx + 3 < n; idx += stride) {
float4 in = *reinterpret_cast<const float4*>(input + idx);
float4 out;
#pragma unroll
for (int i = 0; i < 4; ++i) {
float x = (&in.x)[i];
float y = (x > 20.0f) ? x : log1pf(expf(x));
(&out.x)[i] = y;
}
*reinterpret_cast<float4*>(output + idx) = out;
}
// scalar tail
for (; idx < n; idx += stride) {
float x = input[idx];
float y = (x > 20.0f) ? x : log1pf(expf(x));
output[idx] = y;
}
```
Wait, after the float4 loop, `idx` has been incremented by `stride`. But the tail loop should start from where? Actually, both loops use `idx += stride`, so the tail loop covers elements that the float4 loop skipped because `idx + 3 >= n`. Since `stride` is a multiple of 4? Not necessarily. `stride = gridDim.x * blockDim.x * 4`. If gridDim.x * blockDim.x is not a multiple of something, `idx` in the tail might not cover all elements correctly? Let's think.
If we set `idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4`, then thread 0 handles 0,1,2,3; thread 1 handles 4,5,6,7; etc. The `stride` for the float4 loop is `gridDim.x * blockDim.x * 4`. So each thread processes every `stride`th *group* of 4 elements.
For the tail, we want threads to process the remaining elements individually. But if we keep the same `idx` and `stride`, the tail loop would have `idx` values that are the same as in the float4 loop but with `stride` unchanged. That means thread 0 would process 0, stride, 2*stride... which are already covered in the float4 loop? No, the float4 loop covers 0..3, stride..stride+3, etc. The tail elements are those where `idx` is in `[n-3, n-1]` but not covered by the float4 loop because `idx+3 >= n`. However, these tail elements are not aligned to the start of a float4 group. For example, if n=10, the float4 groups are at 0 and 4. The elements 8 and 9 are not in a complete group. But thread 2 (handles 8..11) would skip the float4 loop because 8+3=11 >= 10. Then in the tail loop, thread 2 would start at idx=8 and process 8, then 8+stride, etc. But what about thread 3 (idx=12)? It would skip both loops. So elements 8 and 9 are handled by thread 2 in the tail loop. But wait, is there any element that is missed? Let's formalize.
If `n` is not a multiple of 4, the last few elements are `n - rem` to `n-1`, where `rem = n % 4`.
The float4 groups start at `0, 4, 8, ..., n - rem`. The last group of 4 is at `n - rem` if `rem == 0`, else the last full group is at `n - rem - 4`? Actually, `n = 4 * q + r`, `0 <= r < 4`. The full groups are at `0, 4, ..., 4*(q-1)`. The element `4*q` is the start of a partial group.
Wait, `4*q` elements are fully covered in q groups. The remaining r elements start at index `4*q`.
In our loop, thread i starts at `idx0 = i * 4`. It processes `idx0, idx0 + stride, idx0 + 2*stride, ...`. If `stride` is a multiple of 4 (which it is, since `stride = num_threads * 4`), then all `idx` values are multiples of 4. So each thread processes groups starting at multiples of 4. The last full group is at index `4*(q-1)`. The partial group starts at `4*q`. Since `4*q` is a multiple of 4, some thread will have `idx = 4*q`. That thread will check `idx + 3 < n` i.e., `4*q + 3 < 4*q + r`, which is `3 < r`. If `r > 3`? No, `r < 4`. So if `r=0`, it's a full group and processed. If `r>0`, then `4*q` is a partial group and skipped by the float4 loop. Then the tail loop needs to process `4*q, 4*q+1, ..., 4*q+r-1`.
But in the tail loop, thread with `idx = 4*q` processes `4*q` and then `4*q + stride`. Since `stride` is a multiple of 4, `4*q + stride` is also a multiple of 4, but it's `> n` if `stride > r` (which it is for large n). Actually, if `stride` is huge (e.g., 3M * 4 = 12M), then the tail loop for that thread only processes one element: `4*q`. But what about `4*q+1`? It needs to be processed by some thread. Which thread has `idx0 = 4*q+1`? None, because all `idx0` are multiples of 4. So the tail loop as written will *not* cover `4*q+1`, `4*q+2`, `4*q+3` because no thread starts at those indices.
This is a classic issue with vectorized kernels. To handle the tail properly, one approach is:
1. Process `n - (n % 4)` elements with the float4 kernel.
2. Process the remaining `n % 4` elements with a scalar kernel.
So we should launch two kernels or handle it differently. But for a tensor of size 1.6B, the tail is at most 3 elements, which is negligible. However, correctness requires handling all elements. Launching a second tiny kernel is fine.
Alternatively, we can have each thread handle a contiguous chunk of elements using a scalar loop, but process the chunk in groups of 4 manually.
Simpler approach: each thread processes a contiguous range of elements. In a grid-stride loop, that's not natural because grid-stride loops give strided access. But we can have each thread handle a contiguous chunk by assigning a start and end index based on thread/block index.
For example:
```cpp
int64_t tid = blockIdx.x * blockDim.x + threadIdx.x;
int64_t num_threads = gridDim.x * blockDim.x;
int64_t chunk_size = (n + num_threads - 1) / num_threads;
int64_t start = tid * chunk_size;
int64_t end = min(start + chunk_size, n);
```
Then each thread processes from `start` to `end`. Within that, we can use float4 for the portion `start` to `end - (end % 4)` and scalar for the tail. This ensures all elements are covered and contiguous access. Coalescing might be slightly worse than strided grid-stride? Actually, with contiguous chunks and coalesced warp access: if thread 0 handles 0..C-1, thread 1 handles C..2C-1, etc., then when executing, the memory accesses within a warp are to contiguous addresses if `chunk_size` is a multiple of warp size. But if `chunk_size` is large (e.g., 500), then thread 0 accesses 0,1,2,... while thread 1 accesses 500,501,... These are far apart, so within a warp, the addresses are not contiguous. That leads to strided access inside a warp, which hurts coalescing. Actually, warp coalescing requires threads in a warp to access adjacent memory locations simultaneously. In the chunk approach, threads in a warp access locations offset by `chunk_size`, which can be many cache lines apart. That's terrible for coalescing.
So grid-stride loop with strided indexing by thread is better for coalescing: thread i accesses `i, i+stride, i+2*stride...`. In a warp (threads 0..31), the first access is to `0,1,2,...31` which is perfectly coalesced. The next access is `stride, stride+1,...` also coalesced. So grid-stride is best.
To use vectorized loads with grid-stride, we need to handle the tail correctly. The easiest way is to process `n / 4` elements with float4 using a grid-stride loop, but re-index so that each thread processes every `stride`th float4. The number of float4 elements is `n4 = n / 4`. Then we launch a scalar kernel for the remaining `n % 4` elements. But `n % 4` can be up to 3. Actually, launching a kernel for 3 elements might have overhead, but it's negligible compared to 1.6B.
Alternatively, we can just do scalar loads. If the reference is also scalar, we might match it or be slightly faster with fewer instructions. Let's consider using half4 or `__half2`? No, it's fp32.
Given that the reference is PyTorch's `softplus`, which is a standard elementwise kernel, I think a custom scalar kernel with `--use_fast_math` might be enough. But let's be thorough and implement a vectorized version with grid-stride for the main body and a separate scalar tail kernel.
Wait, another option: we can use `c10::cuda::CUDAStreamGuard` but that's not necessary.
Let's also consider using `__ldg` for readonly data. `const float* __restrict__` might already allow compiler to use LDG instructions. We could explicitly use `__ldg(&input[idx])` to load through the read-only data cache. On H100, this might help slightly.
Also, using `expf` and `log1pf` from `math.h`. With `--use_fast_math`, these map to faster intrinsic approximations. The tolerance is 1e-4. `fast_math` changes by a small relative error, usually < 1e-4? Actually, `expf` in fast math can be off by a few ULPs, but since we then take `log1p`, it might be okay. However, for large negative x, `expf(x)` underflows to 0, and `log1p(0)` is 0, which is accurate. For x around -10, `exp(-10) ~ 4.5e-5`, `log1p(4.5e-5) ~ 4.5e-5`. The fast math `expf` might have relative error of maybe 1e-3? That would be multiplied by 4.5e-5, giving absolute error ~4.5e-8, well within 1e-4. For x around 0, `exp(0)=1`, `log1p(1)=log(2)~0.693`. An error in `exp` of 1e-3 relative gives `exp=1.001`, `log1p(1.001)=log(2.001)~0.6936`, error is about 0.0001, which is borderline. Actually, `fast_math` `expf` error is documented as up to 2 ULPs or so, not 1e-3 relative. Let me check: CUDA's `__expf` (fast math) has a maximum error of about 2 ULPs? Actually, `--use_fast_math` maps `expf` to `__expf` which has maximum absolute error of about `+/- 1` bit? The documentation says `__expf` intrinsic has a maximum absolute error of 2^-23.2? No. Actually, for `__expf`, the maximum relative error is about 2^-5 = 3%? Wait, that's wrong. The fast math version of `expf` is less accurate but still quite good. Let me recall: `__expf` has maximum absolute error of about 2 ULPs? I think it's better than that. Actually, NVCC fast math functions trade accuracy for speed, but `__expf` error is typically around 1 ULP or a few ULPs. Let me check memory: For `__expf`, the maximum relative error is about 2^-10? No, that also seems high.
Actually, CUDA documentation (from memory): `__expf` has maximum absolute error of `+/- 1`? I don't recall exactly. But PyTorch's ATen likely uses its own vectorized implementation or calls `expf` as well. It might use `std::exp` or `TH_exp`.
To be safe, let's first compile without `--use_fast_math`. If the performance is good, we keep it. If not, we can try with it.
But wait, we can also try a simpler approach: write a Triton kernel? No, the instructions say use `torch.utils.cpp_extension.load_inline` to compile and bind CUDA source. So it must be a CUDA kernel.
Another thought: can we use `torch.cuda.get_device_capability()`? Not needed.
Another optimization: We can use `launch_bounds` to hint the compiler about occupancy. For 512 threads, maximum blocks per SM is 2 (since 1024 threads per SM for some arch? Actually H100 has 2048 threads per SM, so 512 threads can fit 4 blocks). Using `__launch_bounds__(512)` might not be necessary.
Let's also consider using 1024 threads per block. That might be slightly better.
Actually, let's calculate theoretical bandwidth. 6.4 GB read+write. 3.35 TB/s = 3350 GB/s. Time = 6.4 / 3350 = 1.9 ms. We should aim for < 2.5 ms.
PyTorch's reference might already achieve near this. Let's see how much faster we can be.
Potential improvements over PyTorch:
- PyTorch's softplus might not use `log1pf` but `logf(1.0f + expf(x))`. `log1pf` is more accurate and potentially faster because `expf(x)` for large x is huge. But with threshold=20, we avoid that. `log1pf` might be slightly slower or same speed. Actually, `log1pf(expf(x))` is the standard stable implementation. PyTorch might do `x * beta` and then use something like `log1p(exp(x))`. I think it uses `log1p(exp(x))` already. But maybe not. Let's check ATen source `aten/src/ATen/native/cuda/ActivationSoftplusKernel.cu`. I recall PyTorch's softplus uses `gpu_kernel` with a lambda:
```cpp
auto iter = TensorIterator::unary_op(result, self);
gpu_kernel(iter, [beta, threshold] GPU_LAMBDA(scalar_t a) -> scalar_t {
return (a * beta) > threshold ? a : static_cast<scalar_t>(std::log1p(std::exp(a * beta))) / beta;
});
```
This might use `std::log1p` and `std::exp`. On CUDA, `std::exp` for float might map to `expf`, and `std::log1p` might map to `log1pf` (actually `log1p` is C++11). But due to `TensorIterator` overhead (e.g., dynamic dispatch, contiguity checks, splitting into sub-iterators), there might be some overhead. A fused custom kernel on a contiguous tensor avoids `TensorIterator` overhead.
Also, PyTorch's `TensorIterator` might not fully vectorize the load/store for this particular operation because it's a generic framework. Our custom kernel can do explicit vectorized loads.
Let's implement a vectorized float4 kernel with a scalar tail.
Plan:
1. Scalar kernel for tail (max 3 elements, but we can just launch a separate kernel for the remainder).
2. Vectorized kernel for the main part.
Better: we can handle the tail in the same kernel but with a careful loop. Actually, there is a neat trick: process `n4 = n / 4` float4s using a grid-stride loop. The remaining `r = n % 4` elements are at the end. We can launch the scalar kernel starting from `n4 * 4` with `r` elements. This is clean.
Wait, but `n4 * 4` might not be aligned to a specific float4 boundary? `n4 * 4` is divisible by 4, so the next element is at `n4 * 4`. The scalar kernel can start there. Coalescing for 3 elements doesn't matter.
But there is an even simpler approach: just use a grid-stride loop where each thread processes 4 elements at a time, but not via float4 loads, just scalar loads in a `#pragma unroll 4` loop. This reduces loop overhead and the compiler might be able to issue 4 loads together, but it won't combine them into 128-bit transactions automatically. However, memory bandwidth is usually limited by DRAM, not by instruction throughput, for elementwise ops. So scalar loads with a large grid might already be enough.
Let's estimate SMs and throughput. H100 has 132 SMs. Peak HBM bandwidth is 3.35 TB/s. Each SM has high memory bandwidth to L2. The limiting factor is DRAM bandwidth. To saturate it, we need enough memory requests in flight. Each thread can have multiple loads in flight. With 512 threads/block, 4 blocks/SM = 2048 threads/SM, 132 SMs = 270k threads total. Each thread processes maybe 6000 elements (since 1.6B / 270k ≈ 6000). That's plenty. So even with scalar loads, we should saturate memory bandwidth.
The real question is: can we beat PyTorch's reference? Probably yes, because PyTorch's `TensorIterator` has overhead, and our custom kernel avoids it.
Let's write a robust implementation and compile/run it.
Note on `load_inline`: It might cache the compiled extension. We should give it a unique name or use `verbose=True` if needed. But we can just use `name='softplus_ext'`.
One more thing: the kernel must be compiled for the correct architecture. `load_inline` usually handles this by adding `-gencode arch=compute_70,code=sm_70` etc based on the current GPU. For H100 (sm_90), it should add sm_90 automatically. We don't need to pass arch flags.
Let's draft the code.
```python
import torch
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void softplus_kernel(const float* __restrict__ input,
float* __restrict__ output,
int64_t n) {
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
const int64_t stride = gridDim.x * blockDim.x;
for (; idx < n; idx += stride) {
float x = input[idx];
float y = (x > 20.0f) ? x : log1pf(expf(x));
output[idx] = y;
}
}
torch::Tensor softplus_cuda(torch::Tensor input) {
auto output = torch::empty_like(input);
int64_t n = input.numel();
const int threads = 512;
const int blocks = std::min<int64_t>((n + threads - 1) / threads, 65535 * 32);
softplus_kernel<<<blocks, threads>>>(input.data_ptr<float>(),
output.data_ptr<float>(), n);
return output;
}
'''
CPP_SOURCE = r'''
torch::Tensor softplus_cuda(torch::Tensor input);
'''
module = load_inline(
name='softplus_ext',
cpp_sources=CPP_SOURCE,
cuda_sources=CUDA_SOURCE,
functions=['softplus_cuda'],
extra_cuda_cflags=['-O3'],
)
class ModelNew(torch.nn.Module):
def __init__(self):
super(ModelNew, self).__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return module.softplus_cuda(x)
```
I'll compile and test correctness.