53_Gemm_Scaling_Hardtanh_GELU • warp_divergence_minimized_gemm_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>
#define BLOCK_DIM_X 32
#define BLOCK_DIM_Y 32
__device__ inline float fast_tanh(float x) {
// Faster approximation of tanh using rational function
const float x2 = x * x;
return x * (27.0f + x2) / (27.0f + 9.0f * x2);
}
template <typename scalar_t>
__global__ void fused_activation_kernel(
scalar_t* __restrict__ x,
const scalar_t scaling_factor,
const scalar_t hardtanh_min,
const scalar_t hardtanh_max,
const int rows,
const int cols) {
const int tid_x = threadIdx.x;
const int tid_y = threadIdx.y;
const int gid_x = blockIdx.x * BLOCK_DIM_X + tid_x;
const int gid_y = blockIdx.y * BLOCK_DIM_Y + tid_y;
// Constants for GELU
const scalar_t c = (scalar_t)0.044715;
const scalar_t sqrt_2_over_pi = (scalar_t)0.7978845608028654;
scalar_t val = 0;
if (gid_y < rows && gid_x < cols) {
const int idx = gid_y * cols + gid_x;
val = x[idx];
}
// Scaling
val *= scaling_factor;
// Hardtanh
val = fminf(fmaxf(val, hardtanh_min), hardtanh_max);
// GELU
scalar_t x_cube = val * val * val;
scalar_t tanh_arg = sqrt_2_over_pi * (val + c * x_cube);
scalar_t tanh_res = fast_tanh(tanh_arg);
val = 0.5 * val * (1.0 + tanh_res);
if (gid_y < rows && gid_x < cols) {
const int idx = gid_y * cols + gid_x;
x[idx] = val;
}
}
void fused_activation_cuda(
torch::Tensor& x,
double scaling_factor,
double hardtanh_min,
double hardtanh_max) {
const int rows = x.size(0);
const int cols = x.size(1);
dim3 threads(BLOCK_DIM_X, BLOCK_DIM_Y);
dim3 blocks(
(cols + BLOCK_DIM_X - 1) / BLOCK_DIM_X,
(rows + BLOCK_DIM_Y - 1) / BLOCK_DIM_Y
);
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "fused_activation_cuda", ([&] {
fused_activation_kernel<scalar_t><<<blocks, threads>>>(
x.data_ptr<scalar_t>(),
(scalar_t)scaling_factor,
(scalar_t)hardtanh_min,
(scalar_t)hardtanh_max,
rows,
cols);
}));
}
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, "Module function forward (CUDA)");
}| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 0.914 | inst/cycle | 0.000 | 5 |
| Executed Ipc Elapsed | 0.152 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 24.732 | % | 0.087 | 5 |
| Issued Ipc Active | 0.990 | inst/cycle | 0.000 | 5 |
| SM Busy | 24.732 | % | 0.087 | 5 |
| Memory Throughput | 82109170192.490 | byte/second | 3947690789393296384.000 | 5 |
| Mem Busy | 11.412 | % | 0.084 | 5 |
| Max Bandwidth | 7.348 | % | 0.030 | 5 |
| L1/TEX Hit Rate | 50.000 | % | 0.000 | 5 |
| L2 Hit Rate | 82.818 | % | 0.017 | 5 |
| Mem Pipes Busy | 3.498 | % | 0.006 | 5 |
| Warp Cycles Per Issued Instruction | 27.998 | cycle | 0.150 | 5 |
| Warp Cycles Per Executed Instruction | 30.254 | cycle | 0.138 | 5 |
| Avg. Active Threads Per Warp | 32.000 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 30.080 | 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 | 4.000 | block | 0.000 | 5 |
| Block Limit Shared Mem | 8.000 | block | 0.000 | 5 |
| Block Limit Warps | 2.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 | 43.234 | % | 0.173 | 5 |
| Achieved Active Warps Per SM | 27.670 | warp | 0.070 | 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 (43.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::matmul | ||
| CPU Time | 340725.93 | μs |
| Device Time | 140956.91 | μs |
| Self CPU Time | 11612.98 | μ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 | 329112.95 | μs |
| Device Time | 140956.91 | μs |
| Self CPU Time | 193019.34 | μs |
| Self Device Time | 140956.91 | μ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 | 402780.98 | μs |
| Device Time | 14462.29 | μs |
| Self CPU Time | 402780.98 | μs |
| Self Device Time | 14462.29 | μ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 | 357843.50 | μs |
| Device Time | 763622.39 | μs |
| Self CPU Time | 21092.14 | μ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 | 336752.72 | μs |
| Device Time | 763622.39 | μs |
| Self CPU Time | 23152.14 | μs |
| Self Device Time | 763622.39 | μ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 | 763622.39 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 763622.39 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45289 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.