correct
22_Tanh.py
started
2026-05-02T17:47:01Z
timingkernelbench.timing
kernel mean / median / std
4.64 / — / 0.15 μs
ref mean / median / std
4.30 / — / 0.09 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
0.50x
energyNVML (pynvml)
ratio (ref / kernel)
0.67x
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
dram bandwidth
2839.27 GB/s
warp stalls (top 5)nsight
reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Simple model that performs a Tanh activation.
"""
def __init__(self):
super(Model, self).__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Applies Tanh activation to the input tensor.
Args:
x (torch.Tensor): Input tensor of any shape.
Returns:
torch.Tensor: Output tensor with Tanh applied, same shape as input.
"""
return torch.tanh(x)
batch_size = 4096
dim = 393216
def get_inputs():
x = torch.rand(batch_size, dim)
return [x]
def get_init_inputs():
return [] # No special initialization inputs neededturn 07.2s turn · 7.2s LLM · 0.0s tools
reasoning
We need to write a custom CUDA kernel for tanh. A simple elementwise operation tanh(x). There's potential to use vectorized loads/stores, avoid compute-bound being memory-bound, maybe use approximate tanh. Since the input is large: batch_size=4096, dim=393216, total elements = 4096*393216 = 1,610,612,736 ≈ 1.6 billion elements? Let's compute: 393,216 * 4096 = 393,216*4,096. 393,216*4k = 1,572,864,000, plus 393,216*96 = 37,748,736, total 1,610,612,736. About 1.61e9 elements. Each float32 is 4 bytes => ~6.44 GB. That's a lot: ~6.44 GB memory bandwidth if we read x and write y: 3 reads/writes? Actually read x (4 bytes), write y (4 bytes) => 8 bytes per element => ~12.9 GB of memory traffic. For a GPU with maybe 900-1000 GB/s bandwidth, that would take maybe ~14 ms. But let's check GPU specs.
Let's call get_gpu_specs.
call → get_gpu_specs
callget_gpu_specs
{}
result ← get_gpu_specs (ok)
okget_gpu_specs
GPU specs for NVIDIA H100 80GB HBM3:
total memory (runtime): 79.1 GB
spec entry: H100
GPU Architecture: Hopper
GPU Memory: 80GB
Memory Bandwidth: 3.35 TB/s
FP64 TFLOPS: 34
FP64 Tensor Core TFLOPS: 67
FP32 TFLOPS: 67
TF32 Tensor Core TFLOPS: 989 with sparsity
BFLOAT16 Tensore Core TFLOPS: 1979 with sparsity
FP16 Tensor Core TFLOPS: 1979 with sparsity
FP8 Tensor Core TFLOPS: 3958 with sparsity
INT8 Tensor Core TOPS: 3958 with sparsity
Register File Size: 64K 32-bit registers per SM
Maximum number of registers per thread: 255
Maximum number of thread blocks per SM: 32
Shared memory capacity per SM: 228 KB
Maximum shared memory per thread block: 227 KB
turn 1 · 1 tool call(s) · 11.1s turn · 11.1s LLM · 7,005 chars (click to expand)
turn 111.1s turn · 11.1s LLM · 0.0s tools
reasoning
Let me analyze the task:
1. We need to replace `torch.tanh(x)` with a custom CUDA kernel
2. The input is of shape [4096, 393216] which is quite large: 4096 * 393216 = 1,610,612,736 elements (~6.44 GB for fp32)
3. GPU is H100 with 3.35 TB/s memory bandwidth
4. This is a purely element-wise operation, so it's memory bandwidth bound
For tanh, the operation is:
- Read x
- Compute tanh(x)
- Write output
Theoretical memory traffic: 2 reads/writes per element = 8 bytes per element (fp32)
Total memory: 1.6B elements * 8 bytes = ~12.9 GB
At 3.35 TB/s, theoretical minimum time is ~3.8 ms
For a simple element-wise kernel, I should:
1. Write a simple CUDA kernel that computes tanh for each element
2. Use vectorized loads/stores where possible
3. Use sufficient threads and blocks to saturate the GPU
Since tanh has a well-known formula: tanh(x) = (e^x - e^(-x)) / (e^x + e^(-x)), but CUDA has `tanhf`. However, `tanhf` might be slightly slower than necessary.
Actually, let me think about optimization strategies:
1. Simple kernel with float4 vectorized memory access - this will help bandwidth significantly
2. Use `tanhf` from CUDA math library
3. Grid-stride loop pattern for flexibility
Let me write a kernel using float4 for vectorized memory operations.
For float4, each thread processes 4 elements. With 4096 * 393216 = 1,610,612,736 elements, we need 402,653,184 float4 operations. At high occupancy this should be fine.
Let me write the kernel.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "\nimport torch\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void tanh_kernel(const float* __restrict__ input, float* __restrict__ output, int64_t n) {\n // Vectorized loads/stores via float4\n int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n int64_t stride = gridDim.x * blockDim.x;\n \n const float4* in_vec = reinterpret_cast<const float4*>(input);\n float4* out_vec = reinterpret_cast<float4*>(output);\n int64_t n_vec = n >> 2;\n \n for (int64_t i = idx; i < n_vec; i += stride) {\n float4 x = in_vec[i];\n float4 y;\n y.x = tanhf(x.x);\n y.y = tanhf(x.y);\n y.z = tanhf(x.z);\n y.w = tanhf(x.w);\n out_vec[i] = y;\n }\n \n // Handle remaining elements\n int64_t remainder_start = n_vec << 2;\n for (int64_t i = remainder_start + idx; i < n; i += stride) {\n output[i] = tanhf(input[i]);\n }\n}\n\ntorch::Tensor tanh_forward(torch::Tensor input) {\n auto output = torch::empty_like(input);\n int64_t n = input.numel();\n \n const int threads = 256;\n const int blocks = 2048;\n \n tanh_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n output.data_ptr<float>(),\n n\n );\n \n return output;\n}\n'''\n\nmodule = load_inline(\n name='tanh_fast',\n cpp_sources='',\n cuda_sources=cuda_source,\n functions=['tanh_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\nclass ModelNew(torch.nn.Module):\n def __init__(self):\n super(ModelNew, self).__init__()\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return module.tanh_forward(x)\n"
}
result ← compile_kernel (fail)
failcompile_kernel
compile_kernel FAILED: RuntimeError.
RuntimeError: Error building extension 'tanh_fast': [1/3] /opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=tanh_fast -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/include/python3.10 -fPIC -std=c++17 -c /scratch/tejas/PopcornBench/runs/l12_default/original/Kimi-K2.6/level_1_problem_22_cache/k_63f8cc0fd9d8/tanh_fast/main.cpp -o main.o
FAILED: [code=1] main.o
/opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=tanh_fast -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/include/python3.10 -fPIC -std=c++17 -c /scratch/tejas/PopcornBench/runs/l12_default/original/Kimi-K2.6/level_1_problem_22_cache/k_63f8cc0fd9d8/tanh_fast/main.cpp -o main.o
/scratch/tejas/PopcornBench/runs/l12_default/original/Kimi-K2.6/level_1_problem_22_cache/k_63f8cc0fd9d8/tanh_fast/main.cpp: In function 'void pybind11_init_tanh_fast(pybind11::module_&)':
/scratch/tejas/PopcornBench/runs/l12_default/original/Kimi-K2.6/level_1_problem_22_cache/k_63f8cc0fd9d8/tanh_fast/main.cpp:4:51: error: 'tanh_forward' was not declared in this scope
4 | m.def("tanh_forward", torch::wrap_pybind_function(tanh_forward), "tanh_forward");
| ^~~~~~~~~~~~
[2/3] /usr/local/cuda/bin/nvcc -MD -MF cuda.cuda.o.d -ccbin /opt/conda/bin/x86_64-conda-linux-gnu-cc -DTORCH_EXTENSION_NAME=tanh_fast -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/include/python3.10 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_90,code=sm_90 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++17 -c /scratch/tejas/PopcornBench/runs/l12_default/original/Kimi-K2.6/level_1_problem_22_cache/k_63f8cc0fd9d8/tanh_fast/cuda.cu -o cuda.cuda.o
nvcc warning : incompatible redefinition for option 'compiler-bindir', the last value of this option was used
ninja: build stopped: subcommand failed.
turn 26.6s turn · 6.6s LLM · 0.0s tools
reasoning
The issue is that the cpp_sources argument is empty, but `load_inline` generates a pybind11 module from functions list. The function needs to be declared in the cpp_source for the pybind11 module to see it. The problem is that the CUDA source contains the function definition, but the auto-generated main.cpp doesn't have a forward declaration.
When using `load_inline`, the CUDA source is compiled separately, and the generated cpp file references the functions. However, it seems like the function needs to be declared (or at least visible) to the C++ compiler. Maybe I need to provide a header declaration in `cpp_sources`, or put the function in `cpp_sources` and call the kernel from there. Actually, the issue is that the auto-generated binding code needs to see the function declaration. Let me put a forward declaration in `cpp_sources`.
Wait, another approach is to use `load_inline` with the function in `cpp_sources` and call CUDA kernel from there, or provide the forward declaration in `cpp_sources`.
Let me provide a `cpp_sources` with a forward declaration of the torch::Tensor function and then implement it in `cuda_sources`. Wait, but the function needs to be defined somewhere visible to the C++ linker. Since the CUDA source compiles to an object file and links with the cpp source, a forward declaration in cpp_sources should work.
Let me try:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "\nimport torch\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_source = '''\n#include <torch/extension.h>\ntorch::Tensor tanh_forward(torch::Tensor input);\n'''\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void tanh_kernel(const float* __restrict__ input, float* __restrict__ output, int64_t n) {\n int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n int64_t stride = gridDim.x * blockDim.x;\n \n const float4* in_vec = reinterpret_cast<const float4*>(input);\n float4* out_vec = reinterpret_cast<float4*>(output);\n int64_t n_vec = n >> 2;\n \n for (int64_t i = idx; i < n_vec; i += stride) {\n float4 x = in_vec[i];\n float4 y;\n y.x = tanhf(x.x);\n y.y = tanhf(x.y);\n y.z = tanhf(x.z);\n y.w = tanhf(x.w);\n out_vec[i] = y;\n }\n \n int64_t remainder_start = n_vec << 2;\n for (int64_t i = remainder_start + idx; i < n; i += stride) {\n output[i] = tanhf(input[i]);\n }\n}\n\ntorch::Tensor tanh_forward(torch::Tensor input) {\n auto output = torch::empty_like(input);\n int64_t n = input.numel();\n \n const int threads = 256;\n const int blocks = 2048;\n \n tanh_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n output.data_ptr<float>(),\n n\n );\n \n return output;\n}\n'''\n\nmodule = load_inline(\n name='tanh_fast',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['tanh_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\nclass ModelNew(torch.nn.Module):\n def __init__(self):\n super(ModelNew, self).__init__()\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return module.tanh_forward(x)\n"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 34.8s turn · 4.8s LLM · 0.0s tools
reasoning
Great, the new version with `cpp_sources` specified compiled successfully. Now let me run the correctness check.
call → run_correctness
callrun_correctness
{
"kernel_code": "\nimport torch\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_source = '''\n#include <torch/extension.h>\ntorch::Tensor tanh_forward(torch::Tensor input);\n'''\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void tanh_kernel(const float* __restrict__ input, float* __restrict__ output, int64_t n) {\n int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n int64_t stride = gridDim.x * blockDim.x;\n \n const float4* in_vec = reinterpret_cast<const float4*>(input);\n float4* out_vec = reinterpret_cast<float4*>(output);\n int64_t n_vec = n >> 2;\n \n for (int64_t i = idx; i < n_vec; i += stride) {\n float4 x = in_vec[i];\n float4 y;\n y.x = tanhf(x.x);\n y.y = tanhf(x.y);\n y.z = tanhf(x.z);\n y.w = tanhf(x.w);\n out_vec[i] = y;\n }\n \n int64_t remainder_start = n_vec << 2;\n for (int64_t i = remainder_start + idx; i < n; i += stride) {\n output[i] = tanhf(input[i]);\n }\n}\n\ntorch::Tensor tanh_forward(torch::Tensor input) {\n auto output = torch::empty_like(input);\n int64_t n = input.numel();\n \n const int threads = 256;\n const int blocks = 2048;\n \n tanh_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n output.data_ptr<float>(),\n n\n );\n \n return output;\n}\n'''\n\nmodule = load_inline(\n name='tanh_fast',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['tanh_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\nclass ModelNew(torch.nn.Module):\n def __init__(self):\n super(ModelNew, self).__init__()\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return module.tanh_forward(x)\n"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=7.15e-06 mean_abs_err=1.42e-06 max_rel_err=1.12e-05
turn 46.7s turn · 6.7s LLM · 0.0s tools
reasoning
The kernel passes correctness. Let me run the static check and then consider if I want to optimize further before submitting.
Actually, the kernel is pretty simple - it's just tanh. The main optimizations are:
1. Use fast math tanhf
2. Vectorized float4 loads/stores
3. Grid-stride loop for good occupancy
4. 2048 blocks for good throughput on H100
The data size is 4096 * 393216 * 4 bytes ≈ 6.44 GB. At 3.35 TB/s, memory bandwidth bound operations should take ~1.9ms.
With float4 (4 words at a time), we get 4x effective bandwidth, so ~64 bytes per load on H100. This should be very close to memory bandwidth saturation.
Let me do the static check and submit.
call → static_check
callstatic_check
{
"kernel_code": "\nimport torch\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_source = '''\n#include <torch/extension.h>\ntorch::Tensor tanh_forward(torch::Tensor input);\n'''\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void tanh_kernel(const float* __restrict__ input, float* __restrict__ output, int64_t n) {\n int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n int64_t stride = gridDim.x * blockDim.x;\n \n const float4* in_vec = reinterpret_cast<const float4*>(input);\n float4* out_vec = reinterpret_cast<float4*>(output);\n int64_t n_vec = n >> 2;\n \n for (int64_t i = idx; i < n_vec; i += stride) {\n float4 x = in_vec[i];\n float4 y;\n y.x = tanhf(x.x);\n y.y = tanhf(x.y);\n y.z = tanhf(x.z);\n y.w = tanhf(x.w);\n out_vec[i] = y;\n }\n \n int64_t remainder_start = n_vec << 2;\n for (int64_t i = remainder_start + idx; i < n; i += stride) {\n output[i] = tanhf(input[i]);\n }\n}\n\ntorch::Tensor tanh_forward(torch::Tensor input) {\n auto output = torch::empty_like(input);\n int64_t n = input.numel();\n \n const int threads = 256;\n const int blocks = 2048;\n \n tanh_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n output.data_ptr<float>(),\n n\n );\n \n return output;\n}\n'''\n\nmodule = load_inline(\n name='tanh_fast',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['tanh_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\nclass ModelNew(torch.nn.Module):\n def __init__(self):\n super(ModelNew, self).__init__()\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return module.tanh_forward(x)\n"
}
result ← static_check (ok)
okstatic_check
static_check PASSED: no violations or warnings detected.
turn 57.5s turn · 7.5s LLM · 0.0s tools · FINAL
reasoning
The kernel compiled, passed correctness, and the static check is clean. Let me now submit the kernel.
The kernel uses:
- float4 vectorized loads/stores (memory bandwidth optimized)
- A grid-stride loop (allows flexible thread count)
- --use_fast_math for faster tanhf
- 2048 blocks x 256 threads for good occupancy on the GPU
This should provide a nice speedup over PyTorch's native implementation. Let me submit.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "\nimport torch\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_source = '''\n#include <torch/extension.h>\ntorch::Tensor tanh_forward(torch::Tensor input);\n'''\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void tanh_kernel(const float* __restrict__ input, float* __restrict__ output, int64_t n) {\n int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n int64_t stride = gridDim.x * blockDim.x;\n \n const float4* in_vec = reinterpret_cast<const float4*>(input);\n float4* out_vec = reinterpret_cast<float4*>(output);\n int64_t n_vec = n >> 2;\n \n for (int64_t i = idx; i < n_vec; i += stride) {\n float4 x = in_vec[i];\n float4 y;\n y.x = tanhf(x.x);\n y.y = tanhf(x.y);\n y.z = tanhf(x.z);\n y.w = tanhf(x.w);\n out_vec[i] = y;\n }\n \n int64_t remainder_start = n_vec << 2;\n for (int64_t i = remainder_start + idx; i < n; i += stride) {\n output[i] = tanhf(input[i]);\n }\n}\n\ntorch::Tensor tanh_forward(torch::Tensor input) {\n auto output = torch::empty_like(input);\n int64_t n = input.numel();\n \n const int threads = 256;\n const int blocks = 2048;\n \n tanh_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n output.data_ptr<float>(),\n n\n );\n \n return output;\n}\n'''\n\nmodule = load_inline(\n name='tanh_fast',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['tanh_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\nclass ModelNew(torch.nn.Module):\n def __init__(self):\n super(ModelNew, self).__init__()\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return module.tanh_forward(x)\n"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (5 / 5) correctness trials all passed.
Kernel runtime: 4.64 μs
Runtime stats: mean=4.64μs median=?μs std=0.15μs
Numerical precision: max_abs_err=7.15e-06 mean_abs_err=1.42e-06 max_rel_err=1.12e-05
Memory: 12321.0 MB peak (ref: 12321.0 MB, ratio: 1.0x)
Kernel launches: 2 (ref: 1, fusion ratio: 0.5)
Energy: 3669.16 mJ/run (ref: 2471.42 mJ, ratio: 0.6736x)
SOL score: 0.8475
final submitted kernel (71 lines, 1,763 chars) — full source, untruncated
import torch
from torch.utils.cpp_extension import load_inline
cpp_source = '''
#include <torch/extension.h>
torch::Tensor tanh_forward(torch::Tensor input);
'''
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void tanh_kernel(const float* __restrict__ input, float* __restrict__ output, int64_t n) {
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
int64_t stride = gridDim.x * blockDim.x;
const float4* in_vec = reinterpret_cast<const float4*>(input);
float4* out_vec = reinterpret_cast<float4*>(output);
int64_t n_vec = n >> 2;
for (int64_t i = idx; i < n_vec; i += stride) {
float4 x = in_vec[i];
float4 y;
y.x = tanhf(x.x);
y.y = tanhf(x.y);
y.z = tanhf(x.z);
y.w = tanhf(x.w);
out_vec[i] = y;
}
int64_t remainder_start = n_vec << 2;
for (int64_t i = remainder_start + idx; i < n; i += stride) {
output[i] = tanhf(input[i]);
}
}
torch::Tensor tanh_forward(torch::Tensor input) {
auto output = torch::empty_like(input);
int64_t n = input.numel();
const int threads = 256;
const int blocks = 2048;
tanh_kernel<<<blocks, threads>>>(
input.data_ptr<float>(),
output.data_ptr<float>(),
n
);
return output;
}
'''
module = load_inline(
name='tanh_fast',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['tanh_forward'],
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.tanh_forward(x)