← Back to Leaderboard

The AI CUDA Engineer 👷

53_Gemm_Scaling_Hardtanh_GELUshared_mem_optimized_kernel_base

Level 2 • Task 53
import 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]

Kernel Information

Related Kernels (Level 2, Task 53 • 53_Gemm_Scaling_Hardtanh_GELU)

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

constexpr int BLOCK_SIZE = 256;

__device__ inline float scale_value(float val, float factor) {
    return val * factor;
}

__device__ inline float hard_tanh(float val, float min_val, float max_val) {
    return fminf(fmaxf(val, min_val), max_val);
}

__device__ inline float gelu_activation(float val) {
    constexpr float c = 0.044715;
    constexpr float sqrt_2_over_pi = 0.7978845608028654;
    float x_cube = val * val * val;
    float tanh_arg = sqrt_2_over_pi * (val + c * x_cube);
    float tanh_res = tanh(tanh_arg);
    return 0.5 * val * (1.0 + tanh_res);
}

__global__ void fused_activation_with_shared_memory(
    float* __restrict__ x,
    const float scaling_factor,
    const float hardtanh_min,
    const float hardtanh_max,
    const int64_t numel) {
    __shared__ float shared_data[BLOCK_SIZE];

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;

    for (int i = idx; i < numel; i += stride) {
        // Load data into shared memory
        float val = x[i];

        // Use shared memory for subsequent operations
        shared_data[threadIdx.x] = scale_value(val, scaling_factor);
        shared_data[threadIdx.x] = hard_tanh(
            shared_data[threadIdx.x], hardtanh_min, hardtanh_max);
        val = gelu_activation(shared_data[threadIdx.x]);

        // Write back the computed value
        x[i] = 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", ([&] {
        fused_activation_with_shared_memory<<<blocks, BLOCK_SIZE>>>(
            x.data_ptr<float>(),
            static_cast<float>(scaling_factor),
            static_cast<float>(hardtanh_min),
            static_cast<float>(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, "Module function forward (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 0.504 inst/cycle 0.000 5
Executed Ipc Elapsed 0.192 inst/cycle 0.000 5
Issue Slots Busy 13.518 % 0.019 5
Issued Ipc Active 0.542 inst/cycle 0.000 5
SM Busy 13.518 % 0.019 5
Memory Throughput 74711585836.802 byte/second 2680222953560535040.000 5
Mem Busy 10.350 % 0.037 5
Max Bandwidth 6.718 % 0.013 5
L1/TEX Hit Rate 50.000 % 0.000 5
L2 Hit Rate 82.732 % 0.045 5
Mem Pipes Busy 2.218 % 0.002 5
Warp Cycles Per Issued Instruction 25.482 cycle 0.099 5
Warp Cycles Per Executed Instruction 27.324 cycle 0.114 5
Avg. Active Threads Per Warp 30.550 0.000 5
Avg. Not Predicated Off Threads Per Warp 29.160 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 16.000 block 0.000 5
Block Limit Shared Mem 16.000 block 0.000 5
Block Limit Warps 8.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 21.910 % 0.011 5
Achieved Active Warps Per SM 14.022 warp 0.005 5
Analysis Rules
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.0%) 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 637639.96 μs
Device Time 184.28 μs
Self CPU Time 14149.32 μ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::_to_copy
CPU Time 623490.64 μs
Device Time 184.28 μs
Self CPU Time 108.81 μ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::empty_strided
CPU Time 622873.71 μs
Device Time 0.00 μs
Self CPU Time 116.96 μ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
cudaDeviceGetStreamPriorityRange
CPU Time 620626.95 μs
Device Time 0.00 μs
Self CPU Time 620626.95 μ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 331408.72 μs
Device Time 169624.04 μs
Self CPU Time 12403.87 μ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 319004.85 μs
Device Time 169624.04 μs
Self CPU Time 178649.25 μs
Self Device Time 169624.04 μ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 545675.72 μs
Device Time 17437.62 μs
Self CPU Time 545675.72 μs
Self Device Time 17437.62 μ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 496378.41 μs
Device Time 914002.06 μs
Self CPU Time 21134.17 μ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 475245.87 μs
Device Time 914002.06 μs
Self CPU Time 24958.92 μs
Self Device Time 914002.06 μ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 914002.06 μs
Self CPU Time 0.00 μs
Self Device Time 914002.06 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
Status: Completed
45287 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.
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:21:12 bugprone-narrowing-conversions
21 | return 0.5 * val * (1.0 + tanh_res);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:26:5: warning: 2 adjacent parameters of 'fused_activation_with_shared_memory' of similar type ('const float') are easily swapped by mistake [bugprone-easily-swappable-parameters]
26 | const float scaling_factor,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~
27 | const float hardtanh_min,
| ~~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:26:17: note: the first parameter in the range is 'scaling_factor'
26 | const float scaling_factor,
| ^~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:27:17: note: the last parameter in the range is 'hardtanh_min'
27 | const float hardtanh_min,
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:28:5: warning: 2 adjacent parameters of 'fused_activation_with_shared_memory' of convertible types are easily swapped by mistake [bugprone-easily-swappable-parameters]
28 | const float hardtanh_max,
| ^~~~~~~~~~~~~~~~~~~~~~~~~
29 | const int64_t numel) {
| ~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:28:17: note: the first parameter in the range is 'hardtanh_max'
28 | const float hardtanh_max,
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:29:19: note: the last parameter in the range is 'numel'
29 | const int64_t numel) {
| ^~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:28:5: note:
28 | const float hardtanh_max,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:29:5: note: 'const float' and 'const int64_t' may be implicitly converted: 'const float' (as 'float') -> 'const int64_t' (as 'long'), 'const int64_t' (as 'long') -> 'const float' (as 'float')
29 | const int64_t numel) {
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:32:15: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
32 | int idx = blockIdx.x * blockDim.x + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:33:18: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
33 | int stride = gridDim.x * blockDim.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_53/b4_s1_shared_mem_optimized_kernel/base/base.cu:59:5: warning: inside a lambda, '__func__' expands to the name of the function call operator; consider capturing the name of the enclosing function explicitly [bugprone-lambda-function-name]
59 | AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "fused_activation_cuda", ([&] {
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:237:34: note: expanded from macro 'AT_DISPATCH_FLOATING_TYPES'
237 | AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:233:3: note: expanded from macro 'AT_DISPATCH_CASE_FLOATING_TYPES'
233 | AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:74:3: note: expanded from macro 'AT_DISPATCH_CASE'
74 | AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, scalar_t, __VA_ARGS__)
| ^
note: (skipping 1 expansions in backtrace; use -fmacro-backtrace-limit=0 to see all)
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:58:7: note: expanded from macro 'AT_PRIVATE_CHECK_SELECTIVE_BUILD'
58 | AT_ERROR( \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/c10/util/Exception.h:711:32: note: expanded from macro 'AT_ERROR'
711 | C10_EXPAND_MSVC_WORKAROUND(TORCH_CHECK(false, ::c10::str(__VA_ARGS__))); \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/c10/util/Exception.h:536:9: note: expanded from macro 'TORCH_CHECK'
536 | __func__, \
| ^