53_Gemm_Scaling_Hardtanh_GELU • modular_functions_base_edit_1import torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor,
scaling_factor: float,
hardtanh_min: float,
hardtanh_max: float,
weight: torch.Tensor,
bias: torch.Tensor,
) -> torch.Tensor:
"""
Applies GEMM, scaling, hardtanh and GELU activation.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_features)
scaling_factor (float): Factor to scale the GEMM output
hardtanh_min (float): Minimum value for hardtanh
hardtanh_max (float): Maximum value for hardtanh
weight (torch.Tensor): Weight matrix of shape (out_features, in_features)
bias (torch.Tensor): Bias vector of shape (out_features)
Returns:
torch.Tensor: Output tensor after applying GEMM, scaling, hardtanh and GELU,
with shape (batch_size, out_features)
"""
x = F.linear(x, weight, bias)
x = x * scaling_factor
x = F.hardtanh(x, min_val=hardtanh_min, max_val=hardtanh_max)
x = F.gelu(x)
return x
class Model(nn.Module):
"""
Model that performs a GEMM, scaling, hardtanh, and GELU activation.
"""
def __init__(
self, in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max
):
super(Model, self).__init__()
gemm = nn.Linear(in_features, out_features)
self.weight = nn.Parameter(gemm.weight)
self.bias = nn.Parameter(gemm.bias)
self.scaling_factor = scaling_factor
self.hardtanh_min = hardtanh_min
self.hardtanh_max = hardtanh_max
def forward(self, x, fn=module_fn):
return fn(
x,
self.scaling_factor,
self.hardtanh_min,
self.hardtanh_max,
self.weight,
self.bias,
)
batch_size = 128
in_features = 1024
out_features = 512
scaling_factor = 0.5
hardtanh_min = -2
hardtanh_max = 2
def get_inputs():
return [torch.randn(batch_size, in_features)]
def get_init_inputs():
return [in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max]
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Model that performs a GEMM, scaling, hardtanh, and GELU activation.
"""
def __init__(self, in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max):
super(Model, self).__init__()
self.gemm = nn.Linear(in_features, out_features)
self.scaling_factor = scaling_factor
self.hardtanh = nn.Hardtanh(min_val=hardtanh_min, max_val=hardtanh_max)
self.gelu = nn.GELU()
def forward(self, x):
x = self.gemm(x)
x = x * self.scaling_factor
x = self.hardtanh(x)
x = self.gelu(x)
return x
batch_size = 128
in_features = 1024
out_features = 512
scaling_factor = 0.5
hardtanh_min = -2
hardtanh_max = 2
def get_inputs():
return [torch.randn(batch_size, in_features)]
def get_init_inputs():
return [in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max]#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
constexpr int BLOCK_SIZE = 512;
template <typename scalar_t>
__device__ inline scalar_t scale_value(scalar_t val, scalar_t factor) {
return val * factor;
}
template <typename scalar_t>
__device__ inline scalar_t hard_tanh(scalar_t val, scalar_t min_val, scalar_t max_val) {
return fminf(fmaxf(val, min_val), max_val);
}
template <typename scalar_t>
__device__ inline scalar_t gelu_activation(scalar_t val) {
constexpr scalar_t c = (scalar_t)0.044715;
constexpr scalar_t sqrt_2_over_pi = (scalar_t)0.7978845608028654;
scalar_t x_cube = val * val * val;
scalar_t tanh_arg = sqrt_2_over_pi * (val + c * x_cube);
scalar_t tanh_res = tanh(tanh_arg);
return 0.5 * val * (1.0 + tanh_res);
}
template <typename scalar_t>
__global__ void optimized_fused_kernel(
scalar_t* __restrict__ x,
scalar_t scaling_factor,
scalar_t hardtanh_min,
scalar_t hardtanh_max,
int64_t numel) {
int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
int stride = gridDim.x * BLOCK_SIZE;
#pragma unroll 2
for (; idx < numel; idx += stride) {
scalar_t val = x[idx];
val = scale_value(val, scaling_factor);
val = hard_tanh(val, hardtanh_min, hardtanh_max);
x[idx] = gelu_activation(val);
}
}
void fused_activation_cuda(
torch::Tensor& x,
double scaling_factor,
double hardtanh_min,
double hardtanh_max) {
const auto numel = x.numel();
const dim3 blocks((numel + BLOCK_SIZE - 1) / BLOCK_SIZE);
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "fused_activation_cuda", ([&] {
optimized_fused_kernel<scalar_t><<<blocks, BLOCK_SIZE>>>(
x.data_ptr<scalar_t>(),
(scalar_t)scaling_factor,
(scalar_t)hardtanh_min,
(scalar_t)hardtanh_max,
numel);
}));
}
torch::Tensor module_fn_forward(
torch::Tensor x,
double scaling_factor,
double hardtanh_min,
double hardtanh_max,
torch::Tensor weight,
torch::Tensor bias) {
x = x.contiguous().cuda();
weight = weight.contiguous().cuda();
bias = bias.contiguous().cuda();
auto xw = torch::matmul(x, weight.t()) + bias;
fused_activation_cuda(xw, scaling_factor, hardtanh_min, hardtanh_max);
return xw;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &module_fn_forward, "Optimized module function forward (CUDA)");
}| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 0.564 | inst/cycle | 0.000 | 5 |
| Executed Ipc Elapsed | 0.176 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 14.718 | % | 0.127 | 5 |
| Issued Ipc Active | 0.586 | inst/cycle | 0.000 | 5 |
| SM Busy | 14.718 | % | 0.127 | 5 |
| Memory Throughput | 80220455514.488 | byte/second | 4011377984223753728.000 | 5 |
| Mem Busy | 11.108 | % | 0.106 | 5 |
| Max Bandwidth | 7.192 | % | 0.026 | 5 |
| L1/TEX Hit Rate | 50.000 | % | 0.000 | 5 |
| L2 Hit Rate | 82.792 | % | 0.216 | 5 |
| Mem Pipes Busy | 4.944 | % | 0.015 | 5 |
| Warp Cycles Per Issued Instruction | 24.120 | cycle | 0.095 | 5 |
| Warp Cycles Per Executed Instruction | 25.294 | cycle | 0.103 | 5 |
| Avg. Active Threads Per Warp | 30.290 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 29.210 | 0.000 | 5 | |
| Max Active Clusters | 0.000 | cluster | 0.000 | 5 |
| Max Cluster Size | 8.000 | block | 0.000 | 5 |
| Overall GPU Occupancy | 0.000 | % | 0.000 | 5 |
| Cluster Occupancy | 0.000 | % | 0.000 | 5 |
| Block Limit SM | 32.000 | block | 0.000 | 5 |
| Block Limit Registers | 5.000 | block | 0.000 | 5 |
| Block Limit Shared Mem | 16.000 | block | 0.000 | 5 |
| Block Limit Warps | 4.000 | block | 0.000 | 5 |
| Theoretical Active Warps per SM | 64.000 | warp | 0.000 | 5 |
| Theoretical Occupancy | 100.000 | % | 0.000 | 5 |
| Achieved Occupancy | 22.596 | % | 0.077 | 5 |
| Achieved Active Warps Per SM | 14.458 | warp | 0.031 | 5 |
| Rule | Description |
|---|---|
| WRN HighPipeUtilization | All compute pipelines are under-utilized. Either this kernel is very small or it doesn't issue enough warps per scheduler. Check the Launch Statistics and Scheduler Statistics sections for further details. |
| INF CPIStall | Check the Warp Stall Sampling (All Cycles) table for the top stall locations in your source based on sampling data. The Kernel Profiling Guide (https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#metrics-reference) provides more details on each stall reason. |
| WRN Occupancy | This kernel's theoretical occupancy is not impacted by any block limit. The difference between calculated theoretical (100.0%) and measured achieved occupancy (22.6%) can be the result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can occur between warps within a block as well as across blocks of the same kernel. See the CUDA Best Practices Guide (https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#occupancy) for more details on optimizing occupancy. |
| Operation / Metric | Value | Unit |
|---|---|---|
| aten::to | ||
| CPU Time | 346703.14 | μs |
| Device Time | 162.69 | μs |
| Self CPU Time | 12619.79 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::matmul | ||
| CPU Time | 337995.77 | μs |
| Device Time | 165724.13 | μs |
| Self CPU Time | 12610.22 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::mm | ||
| CPU Time | 325385.56 | μs |
| Device Time | 165724.13 | μs |
| Self CPU Time | 191485.03 | μs |
| Self Device Time | 165724.13 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| cudaLaunchKernel | ||
| CPU Time | 532592.61 | μs |
| Device Time | 17039.96 | μs |
| Self CPU Time | 532592.61 | μs |
| Self Device Time | 17039.96 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::zero_ | ||
| CPU Time | 489304.76 | μs |
| Device Time | 895246.46 | μs |
| Self CPU Time | 22282.28 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::fill_ | ||
| CPU Time | 467026.79 | μs |
| Device Time | 895246.46 | μs |
| Self CPU Time | 24717.42 | μs |
| Self Device Time | 895246.46 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<int>, at::detail::Array<char*, 1> >(int, at::native::FillFunctor<int>, at::detail::Array<char*, 1>) | ||
| CPU Time | 0.00 | μs |
| Device Time | 895246.46 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 895246.46 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45285 warnings generated when compiling for host. Suppressed 45326 warnings (45279 in non-user code, 47 NOLINT). Use -header-filter=.* to display errors from all non-system headers. Use -system-headers to display errors from system headers as well.