correct
70_Gemm_Sigmoid_Scaling_ResidualAdd.py
started
2026-05-04T04:23:12Z
timingkernelbench.timing
kernel mean / median / std
2.68 / — / 0.02 μs
ref mean / median / std
2.74 / — / 0.01 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
2.00x
energyNVML (pynvml)
ratio (ref / kernel)
0.50x
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
dram bandwidth
2334.72 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):
"""
Model implementing the pattern "Gemm_Sigmoid_Scaling_ResidualAdd".
"""
def __init__(self, input_size, hidden_size, scaling_factor):
super(Model, self).__init__()
self.gemm = nn.Linear(input_size, hidden_size)
self.scaling_factor = scaling_factor
def forward(self, x):
"""
Forward pass of the model.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, input_size).
Returns:
torch.Tensor: Output tensor of shape (batch_size, hidden_size).
"""
x = self.gemm(x)
original_x = x
x = torch.sigmoid(x)
x = x * self.scaling_factor
x = x + original_x
return x
batch_size = 1024
input_size = 8192
hidden_size = 8192
scaling_factor = 2.0
def get_inputs():
return [torch.rand(batch_size, input_size)]
def get_init_inputs():
return [input_size, hidden_size, scaling_factor]turn 08.3s turn · 8.3s LLM · 0.0s tools
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 114.6s turn · 14.6s LLM · 0.0s tools
call → compile_kernel
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 24.2s turn · 4.2s LLM · 0.0s tools
call → run_correctness
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=4.77e-07 mean_abs_err=1.82e-08 max_rel_err=7.69e-02
turn 34.9s turn · 4.9s LLM · 0.0s tools
reasoning
**Evaluating performance needs**
I realize there's a need to assess performance. I should consider profiling to gather data on how well things are working. I wonder what specific metrics would be most helpful. Maybe looking into areas where improvements can be made? It’s important to ensure everything runs smoothly. I’m curious if I should prioritize specific aspects of performance over others. Alright, I’ll start figuring out what would be most effective to profile!
call → profile_kernel
result ← profile_kernel (ok)
okprofile_kernel
profile_kernel PASSED: profiling complete.
=== Kernel Profile Summary ===
--- Kernel Breakdown (by GPU time) ---
#1: aten::addmm 72519.0μs (100%) calls=1
#2: aten::t 101.0μs (0%) calls=1
#3: aten::transpose 44.0μs (0%) calls=1
#4: aten::as_strided 15.0μs (0%) calls=1
--- Timing ---
GPU time: 2721.98 μs
--- Memory Bandwidth ---
DRAM BW: 1480.3/3350 GB/s (44.2%)
read=1138.5 write=341.8 GB/s
Hit rates: L2=140.9%
--- Compute Throughput ---
Dominant pipe: FP32 util=102.2%
Pipe util: FMA=102.2%, Tensor=1.6%
--- Roofline Analysis ---
Ridge point (FP32): 20.00 FLOPs/byte
Classification: COMPUTE-BOUND
--- Occupancy ---
Achieved: 91.9%
Eligible warps/cycle: 3.46
Launch: regs/thread=278, smem/blk=35840B, block=384, grid=33280
--- Memory Access Quality ---
Global ld coalescing: 8.00 sectors/req (8.0x wasted BW)
--- Warp Stall Reasons (% of issue cycles) ---
long_scoreboard 931.1% ####################
short_scoreboard 159.8% ####################
mio_throttle 18.1% ###
--- Optimization Hints ---
-> 931% stalls: global memory latency (waiting for DRAM/L2). + sectors/req=8.0 -> uncoalesced global loads, restructure access pattern or use __ldg Try shared-memory tiling, prefetching, or reducing working set.
-> Global load coalescing=8.0x (ideal=1.0) -> 8x wasted DRAM BW. Make per-warp accesses contiguous.
-> Compute-bound but NOT using tensor cores. If applicable, use FP16/BF16 with torch.matmul or WMMA intrinsics.
turn 416.6s turn · 16.6s LLM · 0.0s tools
reasoning
**Optimizing performance**
I think I need to evaluate whether the timing was included in my last submission. It's important to consider whether I have the complete profile with all the necessary components. Perhaps submitting isn't the best option just yet. It might be worthwhile to optimize through element-wise operations, like vectorized float4 calculations, and maybe even analyze using an approximate sigmoid polynomial. This current kernel using exp seems to take about 100 microseconds, so there's potential for improvement here.
call → run_correctness
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=4.77e-07 mean_abs_err=1.82e-08 max_rel_err=7.69e-02
turn 510.4s turn · 10.4s LLM · 0.0s tools
call → profile_kernel
result ← profile_kernel (ok)
okprofile_kernel
profile_kernel PASSED: profiling complete.
=== Kernel Profile Summary ===
--- Kernel Breakdown (by GPU time) ---
#1: aten::mm 71436.0μs (100%) calls=1
#2: aten::t 95.0μs (0%) calls=1
#3: aten::transpose 38.0μs (0%) calls=1
#4: aten::as_strided 11.0μs (0%) calls=1
--- Timing ---
GPU time: 2769.34 μs [↑47.4μs REGRESSED]
--- Memory Bandwidth ---
DRAM BW: 1242.0/3350 GB/s (37.1%) [↓7.1% REGRESSED]
read=960.4 write=281.6 GB/s
Hit rates: L2=144.0%
--- Compute Throughput ---
Dominant pipe: FP32 util=104.3% [↑2.2% ok]
Pipe util: FMA=104.3%, Tensor=1.3%
--- Roofline Analysis ---
Ridge point (FP32): 20.00 FLOPs/byte
Classification: COMPUTE-BOUND
--- Occupancy ---
Achieved: 95.1% [↑3.2% ok]
Eligible warps/cycle: 4.78
Launch: regs/thread=285, smem/blk=35328B, block=384, grid=33280
--- Memory Access Quality ---
Global ld coalescing: 8.00 sectors/req (8.0x wasted BW)
--- Warp Stall Reasons (% of issue cycles) ---
long_scoreboard 604.1% ####################
short_scoreboard 127.5% ####################
mio_throttle 21.2% ####
--- Optimization Hints ---
-> 604% stalls: global memory latency (waiting for DRAM/L2). + sectors/req=8.0 -> uncoalesced global loads, restructure access pattern or use __ldg Try shared-memory tiling, prefetching, or reducing working set.
-> Global load coalescing=8.0x (ideal=1.0) -> 8x wasted DRAM BW. Make per-warp accesses contiguous.
-> Compute-bound but NOT using tensor cores. If applicable, use FP16/BF16 with torch.matmul or WMMA intrinsics.
--- Delta vs Previous Iteration ---
GPU time: 2721.98→2769.34μs (↑47.36) REGRESSED
DRAM BW%: 44.19→37.07% (↓7.11) REGRESSED
Compute%: 102.19→104.35% (↑2.16) ok
Occupancy: 91.88→95.06% (↑3.18) ok
Eligible warps/cyc: 3.46→4.78 (↑1.32) ok
L2 hit%: 140.92→144.00% (↑3.08) ok
turn 614.8s turn · 14.8s LLM · 0.0s tools
reasoning
**Exploring kernel optimizations**
I’m considering how to submit variants for optimization without rushing. Maybe I can inspect custom elements using static or disassembly. If I set a scaling factor of 2, that could simplify things, especially since the constructor accepts any value but only uses 2.0 in the init function. I wonder if I should store the scale, or if evaluating it uses just the given kernel. The output should have a minimal error margin; perhaps using `torch.sigmoid(x)` for accuracy could help!
**Optimizing float4 kernels**
I’m thinking about how to vectorize elementwise float4 to cut down on indexing overhead and boost memory efficiency. Currently, each thread processes a single scalar. If I use a float4 kernel and ensure that n is divisible by 4 (like 8M), I can load float4 blocks and compute four exponential functions at once. This would reduce memory instructions while keeping the same exponential calculation intact, which sounds like a solid optimization!
call → run_correctness
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=4.77e-07 mean_abs_err=1.82e-08 max_rel_err=7.69e-02
turn 77.3s turn · 7.3s LLM · 0.0s tools
call → profile_kernel
result ← profile_kernel (ok)
okprofile_kernel
profile_kernel PASSED: profiling complete.
=== Kernel Profile Summary ===
--- Kernel Breakdown (by GPU time) ---
#1: aten::addmm 103807.0μs (100%) calls=1
#2: aten::t 105.0μs (0%) calls=1
#3: aten::transpose 43.0μs (0%) calls=1
#4: aten::as_strided 12.0μs (0%) calls=1
--- Timing ---
GPU time: 2705.47 μs [↓63.9μs ok]
--- Memory Bandwidth ---
DRAM BW: 2338.5/3350 GB/s (69.8%) [↑32.7% ok]
read=1774.2 write=564.3 GB/s
Hit rates: L2=140.3%
--- Compute Throughput ---
Dominant pipe: FP32 util=89.4% [↓14.9% REGRESSED]
Pipe util: FMA=89.4%, Tensor=2.1%
--- Roofline Analysis ---
Ridge point (FP32): 20.00 FLOPs/byte
Classification: COMPUTE-BOUND
--- Occupancy ---
Achieved: 90.3% [↓4.8% REGRESSED]
Eligible warps/cycle: 2.30
Launch: regs/thread=283, smem/blk=35840B, block=384, grid=8704
--- Memory Access Quality ---
Global ld coalescing: 20.00 sectors/req (20.0x wasted BW)
--- Warp Stall Reasons (% of issue cycles) ---
long_scoreboard 2716.7% ####################
short_scoreboard 119.4% ####################
mio_throttle 80.5% ################
--- Optimization Hints ---
-> 2717% stalls: global memory latency (waiting for DRAM/L2). + sectors/req=20.0 -> uncoalesced global loads, restructure access pattern or use __ldg Try shared-memory tiling, prefetching, or reducing working set.
-> Global load coalescing=20.0x (ideal=1.0) -> 20x wasted DRAM BW. Make per-warp accesses contiguous.
-> Compute-bound but NOT using tensor cores. If applicable, use FP16/BF16 with torch.matmul or WMMA intrinsics.
--- Delta vs Previous Iteration ---
GPU time: 2769.34→2705.47μs (↓63.87) ok
DRAM BW%: 37.07→69.80% (↑32.73) ok
Compute%: 104.35→89.42% (↓14.93) REGRESSED
Occupancy: 95.06→90.26% (↓4.81) REGRESSED
Eligible warps/cyc: 4.78→2.30 (↓2.48) REGRESSED
Ld sectors/req: 8.00→20.00 (↑12.00) REGRESSED
L2 hit%: 144.00→140.33% (↓3.67) REGRESSED
turn 823.7s turn · 23.7s LLM · 0.0s tools
reasoning
**Evaluating matrix transposition**
I'm working with the equation Y^T = W * X^T, where W is represented in a row-major format as an N x K matrix. I need to consider that the memory representation of W in column-major format would actually translate to W_col = W^T, making it a K x N matrix instead of just W. I want to make sure I have this correctly represented in both formats!
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale);\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <ATen/cuda/CUDABlas.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cublas_v2.h>\n\n__device__ __forceinline__ float fsigmoid(float z) { return 1.0f / (1.0f + __expf(-z)); }\n\n__global__ void bias_sigmoid_residual_kernel4(float* __restrict__ y, const float* __restrict__ bias, long n4, int cols4, float scale) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n long stride = (long)blockDim.x * gridDim.x;\n float4* __restrict__ y4 = reinterpret_cast<float4*>(y);\n const float4* __restrict__ b4 = reinterpret_cast<const float4*>(bias);\n for (long i = idx; i < n4; i += stride) {\n float4 v = y4[i];\n float4 b = b4[(int)(i % cols4)];\n v.x += b.x; v.y += b.y; v.z += b.z; v.w += b.w;\n v.x = v.x + scale * fsigmoid(v.x);\n v.y = v.y + scale * fsigmoid(v.y);\n v.z = v.z + scale * fsigmoid(v.z);\n v.w = v.w + scale * fsigmoid(v.w);\n y4[i] = v;\n }\n}\n\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale) {\n const int64_t M = x.size(0);\n const int64_t K = x.size(1);\n const int64_t N = weight.size(0);\n auto y = torch::empty({M, N}, x.options());\n\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n const float alpha = 1.0f;\n const float beta = 0.0f;\n // Row-major Y = X * weight^T is column-major Y^T = weight * X^T.\n cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N,\n (int)N, (int)M, (int)K,\n &alpha,\n weight.data_ptr<float>(), (int)K,\n x.data_ptr<float>(), (int)K,\n &beta,\n y.data_ptr<float>(), (int)N);\n\n long n = y.numel();\n if (n > 0) {\n int threads = 256;\n long n4 = n >> 2;\n int blocks = (int)((n4 + threads - 1) / threads);\n if (blocks > 65535) blocks = 65535;\n bias_sigmoid_residual_kernel4<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(\n y.data_ptr<float>(), bias.data_ptr<float>(), n4, (int)(N >> 2), (float)scale);\n }\n return y;\n}\n'''\n\n_ext = load_inline(\n name=\"gemm_sigmoid_scale_residual_ext_v4\",\n cpp_sources=cpp_src,\n cuda_sources=cuda_src,\n functions=[\"fused_forward\"],\n with_cuda=True,\n extra_cuda_cflags=[\"-O3\", \"--use_fast_math\"],\n extra_cflags=[\"-O3\"],\n verbose=False,\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, input_size, hidden_size, scaling_factor):\n super(ModelNew, self).__init__()\n self.gemm = nn.Linear(input_size, hidden_size)\n self.scaling_factor = float(scaling_factor)\n\n def forward(self, x):\n return _ext.fused_forward(x, self.gemm.weight, self.gemm.bias, self.scaling_factor)\n"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=4.77e-07 mean_abs_err=1.82e-08 max_rel_err=7.69e-02
turn 97.4s turn · 7.4s LLM · 0.0s tools
call → profile_kernel
callprofile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale);\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <ATen/cuda/CUDABlas.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cublas_v2.h>\n\n__device__ __forceinline__ float fsigmoid(float z) { return 1.0f / (1.0f + __expf(-z)); }\n\n__global__ void bias_sigmoid_residual_kernel4(float* __restrict__ y, const float* __restrict__ bias, long n4, int cols4, float scale) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n long stride = (long)blockDim.x * gridDim.x;\n float4* __restrict__ y4 = reinterpret_cast<float4*>(y);\n const float4* __restrict__ b4 = reinterpret_cast<const float4*>(bias);\n for (long i = idx; i < n4; i += stride) {\n float4 v = y4[i];\n float4 b = b4[(int)(i % cols4)];\n v.x += b.x; v.y += b.y; v.z += b.z; v.w += b.w;\n v.x = v.x + scale * fsigmoid(v.x);\n v.y = v.y + scale * fsigmoid(v.y);\n v.z = v.z + scale * fsigmoid(v.z);\n v.w = v.w + scale * fsigmoid(v.w);\n y4[i] = v;\n }\n}\n\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale) {\n const int64_t M = x.size(0);\n const int64_t K = x.size(1);\n const int64_t N = weight.size(0);\n auto y = torch::empty({M, N}, x.options());\n\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N,\n (int)N, (int)M, (int)K,\n &alpha,\n weight.data_ptr<float>(), (int)K,\n x.data_ptr<float>(), (int)K,\n &beta,\n y.data_ptr<float>(), (int)N);\n\n long n = y.numel();\n if (n > 0) {\n int threads = 256;\n long n4 = n >> 2;\n int blocks = (int)((n4 + threads - 1) / threads);\n if (blocks > 65535) blocks = 65535;\n bias_sigmoid_residual_kernel4<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(\n y.data_ptr<float>(), bias.data_ptr<float>(), n4, (int)(N >> 2), (float)scale);\n }\n return y;\n}\n'''\n\n_ext = load_inline(\n name=\"gemm_sigmoid_scale_residual_ext_v4\",\n cpp_sources=cpp_src,\n cuda_sources=cuda_src,\n functions=[\"fused_forward\"],\n with_cuda=True,\n extra_cuda_cflags=[\"-O3\", \"--use_fast_math\"],\n extra_cflags=[\"-O3\"],\n verbose=False,\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, input_size, hidden_size, scaling_factor):\n super(ModelNew, self).__init__()\n self.gemm = nn.Linear(input_size, hidden_size)\n self.scaling_factor = float(scaling_factor)\n\n def forward(self, x):\n return _ext.fused_forward(x, self.gemm.weight, self.gemm.bias, self.scaling_factor)\n"
}
result ← profile_kernel (ok)
okprofile_kernel
profile_kernel PASSED: profiling complete.
=== Kernel Profile Summary ===
--- Kernel Breakdown (by GPU time) ---
#1: aten::empty 277.0μs (100%) calls=1
--- Timing ---
GPU time: 2751.46 μs [↑46.0μs REGRESSED]
--- Memory Bandwidth ---
DRAM BW: 2345.6/3350 GB/s (70.0%) [↑0.2% ok]
read=1786.8 write=558.8 GB/s
Hit rates: L2=145.3%
--- Compute Throughput ---
Dominant pipe: FP32 util=93.1% [↑3.7% ok]
Pipe util: FMA=93.1%, Tensor=1.4%
--- Roofline Analysis ---
Ridge point (FP32): 20.00 FLOPs/byte
Classification: COMPUTE-BOUND
--- Occupancy ---
Achieved: 90.2% [↓0.0% REGRESSED]
Eligible warps/cycle: 2.77
Launch: regs/thread=285, smem/blk=35328B, block=384, grid=8704
--- Memory Access Quality ---
Global ld coalescing: 20.00 sectors/req (20.0x wasted BW)
--- Warp Stall Reasons (% of issue cycles) ---
long_scoreboard 1849.2% ####################
short_scoreboard 118.7% ####################
mio_throttle 93.7% ##################
--- Optimization Hints ---
-> 1849% stalls: global memory latency (waiting for DRAM/L2). + sectors/req=20.0 -> uncoalesced global loads, restructure access pattern or use __ldg Try shared-memory tiling, prefetching, or reducing working set.
-> Global load coalescing=20.0x (ideal=1.0) -> 20x wasted DRAM BW. Make per-warp accesses contiguous.
-> Compute-bound but NOT using tensor cores. If applicable, use FP16/BF16 with torch.matmul or WMMA intrinsics.
--- Delta vs Previous Iteration ---
GPU time: 2705.47→2751.46μs (↑45.98) REGRESSED
DRAM BW%: 69.80→70.02% (↑0.21) ok
Compute%: 89.42→93.09% (↑3.68) ok
Occupancy: 90.26→90.21% (↓0.05) REGRESSED
Eligible warps/cyc: 2.30→2.77 (↑0.47) ok
L2 hit%: 140.33→145.27% (↑4.94) ok
turn 1011.9s turn · 11.9s LLM · 0.0s tools
call → static_check
result ← static_check (ok)
okstatic_check
static_check PASSED: no violations or warnings detected.
turn 119.8s turn · 9.8s LLM · 0.0s tools · FINAL
call → submit_kernel
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (5 / 5) correctness trials all passed.
Kernel runtime: 2.68 μs
Runtime stats: mean=2.68μs median=?μs std=0.02μs
Numerical precision: max_abs_err=4.77e-07 mean_abs_err=1.82e-08 max_rel_err=7.69e-02
Memory: 609.06 MB peak (ref: 673.06 MB, ratio: 0.9049x)
Kernel launches: 4 (ref: 8, fusion ratio: 2.0)
Energy: 483.22 mJ/run (ref: 241.78 mJ, ratio: 0.5004x)
SOL score: 0.8951
final submitted kernel (67 lines, 2,128 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cpp_src = r'''
#include <torch/extension.h>
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale);
'''
cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
__device__ __forceinline__ float fsigmoid(float z) {
return 1.0f / (1.0f + __expf(-z));
}
__global__ void sigmoid_residual_kernel4(float* __restrict__ y, long n4, float scale) {
long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
long stride = (long)blockDim.x * gridDim.x;
float4* __restrict__ y4 = reinterpret_cast<float4*>(y);
for (long i = idx; i < n4; i += stride) {
float4 v = y4[i];
v.x = v.x + scale * fsigmoid(v.x);
v.y = v.y + scale * fsigmoid(v.y);
v.z = v.z + scale * fsigmoid(v.z);
v.w = v.w + scale * fsigmoid(v.w);
y4[i] = v;
}
}
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale) {
torch::Tensor y = torch::addmm(bias, x, weight.t());
long n = y.numel();
if (n > 0) {
int threads = 256;
long n4 = n >> 2;
int blocks = (int)((n4 + threads - 1) / threads);
if (blocks > 65535) blocks = 65535;
sigmoid_residual_kernel4<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(), n4, (float)scale);
}
return y;
}
'''
_ext = load_inline(
name="gemm_sigmoid_scale_residual_ext_v3",
cpp_sources=cpp_src,
cuda_sources=cuda_src,
functions=["fused_forward"],
with_cuda=True,
extra_cuda_cflags=["-O3", "--use_fast_math"],
extra_cflags=["-O3"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, input_size, hidden_size, scaling_factor):
super(ModelNew, self).__init__()
self.gemm = nn.Linear(input_size, hidden_size)
self.scaling_factor = float(scaling_factor)
def forward(self, x):
return _ext.fused_forward(x, self.gemm.weight, self.gemm.bias, self.scaling_factor)