← Back to Leaderboard

The AI CUDA Engineer 👷

76_Gemm_Add_ReLUhybrid_optimized_kernel_base

Level 2 • Task 76
import torch
import torch.nn as nn
import torch.nn.functional as F


def module_fn(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
) -> torch.Tensor:
    """
    Performs matrix multiplication, adds bias, and applies ReLU activation.

    Args:
        x (torch.Tensor): Input tensor with shape (batch_size, in_features)
        weight (torch.Tensor): Weight matrix with shape (out_features, in_features)
        bias (torch.Tensor): Bias tensor with shape (out_features,)

    Returns:
        torch.Tensor: Output tensor with shape (batch_size, out_features)
    """
    x = F.linear(x, weight)
    x = x + bias
    x = F.relu(x)
    return x


class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, adds a bias term, and applies ReLU.
    """

    def __init__(self, in_features, out_features, bias_shape):
        super(Model, self).__init__()
        gemm = nn.Linear(in_features, out_features, bias=False)
        self.weight = nn.Parameter(gemm.weight)
        self.bias = nn.Parameter(torch.randn(bias_shape) * 0.02)

    def forward(self, x, fn=module_fn):
        return fn(x, self.weight, self.bias)


batch_size = 128
in_features = 1024
out_features = 512
bias_shape = (out_features,)


def get_inputs():
    return [torch.randn(batch_size, in_features)]


def get_init_inputs():
    return [in_features, out_features, bias_shape]
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, adds a bias term, and applies ReLU.
    """
    def __init__(self, in_features, out_features, bias_shape):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=False)
        self.bias = nn.Parameter(torch.randn(bias_shape)*0.02)

    def forward(self, x):   
        """
        Args:
            x (torch.Tensor): Input tensor with shape (batch_size, in_features).
        Returns:
            torch.Tensor: Output tensor with shape (batch_size, out_features).
        """
        x = self.gemm(x)
        x = x + self.bias
        x = torch.relu(x)
        return x

batch_size = 128
in_features = 1024
out_features = 512
bias_shape = (out_features,)

def get_inputs():
    return [torch.randn(batch_size, in_features)]

def get_init_inputs():
    return [in_features, out_features, bias_shape]

Kernel Information

Related Kernels (Level 2, Task 76 • 76_Gemm_Add_ReLU)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 shared_warp_tile_kernel_base 0.03 0.93 1.54
🥇 combined_warp_tile_base 0.03 0.93 1.54
🥉 optimized_block_size_kernel_base 0.03 0.89 1.49
4 warp_tile_ldg_base 0.03 0.87 1.44
4 even_workload_dist_base_base 0.03 0.87 1.44
4 hybrid_warp_tile_kernel_base 0.03 0.87 1.44
4 warp_tile_hybrid_base 0.03 0.87 1.44
8 warp_tile_ldg_opt_base 0.03 0.81 1.36
8 warp_reduction_optimized_base_base 0.03 0.81 1.36
10 optimized_shared_memory_base_base 0.03 0.79 1.32
10 warp_tile_base_base 0.03 0.79 1.32
12 hybrid_optimized_kernel_base 0.04 0.77 1.28
13 warp_reduction_gemm_base 0.04 0.71 1.18
13 warp_tile_aligned_base_base 0.04 0.71 1.18
15 vectorized_warp_unroll_base_base 0.04 0.69 1.15
15 vectorized_warp_unroll_base_edit_1 0.04 0.69 1.15
15 warp_reduction_unrolled_gemm_edit_1 0.04 0.69 1.15
18 unrolled_warp_gemm_edit_1 0.04 0.67 1.12
18 unrolled_warp_gemm_base 0.04 0.67 1.12
18 vectorized_warp_reduction_base 0.04 0.67 1.12
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAGuard.h>

#define WARP_SIZE 32
#define TILE_SIZE 4

__global__ void hybrid_optimized_kernel(const float* __restrict__ x,
                                         const float* __restrict__ weight,
                                         const float* __restrict__ bias,
                                         float* __restrict__ out,
                                         int batch_size,
                                         int in_features,
                                         int out_features) {
    int batch_idx = blockIdx.x;
    int warp_id = threadIdx.x / WARP_SIZE;
    int lane_id = threadIdx.x % WARP_SIZE;
    int warps_per_block = blockDim.x / WARP_SIZE;
    int base_out_group = blockIdx.y * (warps_per_block * TILE_SIZE);
    int out_base = base_out_group + warp_id * TILE_SIZE;

    if (out_base >= out_features) return;

    float sums[TILE_SIZE] = {0.0f, 0.0f, 0.0f, 0.0f};
    const float* x_row = x + batch_idx * in_features;

    for (int tile = 0; tile < TILE_SIZE; tile++) {
        int current_out = out_base + tile;
        if (current_out < out_features) {
            const float* w_row = weight + current_out * in_features;
            int vec_count = in_features / 4;
            const float4* x_vec = reinterpret_cast<const float4*>(x_row);
            const float4* w_vec = reinterpret_cast<const float4*>(w_row);

            #pragma unroll
            for (int i = lane_id; i < vec_count; i += WARP_SIZE) {
                float4 xv = __ldg(x_vec + i);
                float4 wv = __ldg(w_vec + i);
                sums[tile] += xv.x * wv.x + xv.y * wv.y + xv.z * wv.z + xv.w * wv.w;
            }

            int rem_start = vec_count * 4;
            for (int i = rem_start + lane_id; i < in_features; i += WARP_SIZE) {
                sums[tile] += __ldg(x_row + i) * __ldg(w_row + i);
            }
        }
    }

    for (int tile = 0; tile < TILE_SIZE; tile++) {
        for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {
            sums[tile] += __shfl_down_sync(0xffffffff, sums[tile], offset);
        }
    }

    if (lane_id == 0) {
        for (int tile = 0; tile < TILE_SIZE; tile++) {
            int current_out = out_base + tile;
            if (current_out < out_features) {
                float result = sums[tile] + __ldg(bias + current_out);
                out[batch_idx * out_features + current_out] = (result > 0.0f) ? result : 0.0f;
            }
        }
    }
}

torch::Tensor hybrid_optimized_forward(torch::Tensor x,
                                       torch::Tensor weight,
                                       torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");

    int batch_size = x.size(0);
    int in_features = x.size(1);
    int out_features = weight.size(0);

    auto out = torch::empty({batch_size, out_features}, x.options());

    int warps_per_block = 8;
    int threads_per_block = warps_per_block * WARP_SIZE;
    int blocks_y = (out_features + (warps_per_block * TILE_SIZE) - 1) / (warps_per_block * TILE_SIZE);

    dim3 grid(batch_size, blocks_y);
    dim3 block(threads_per_block);

    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    hybrid_optimized_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size,
        in_features,
        out_features
    );

    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &hybrid_optimized_forward, "Optimized GEMM with bias and ReLU using hybrid approach (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 1.372 inst/cycle 0.000 5
Executed Ipc Elapsed 1.262 inst/cycle 0.000 5
Issue Slots Busy 34.384 % 0.062 5
Issued Ipc Active 1.374 inst/cycle 0.000 5
SM Busy 34.384 % 0.062 5
Memory Throughput 74648583234.880 byte/second 51933681649211696.000 5
Mem Busy 63.868 % 0.042 5
Max Bandwidth 62.834 % 0.041 5
L1/TEX Hit Rate 56.906 % 0.328 5
L2 Hit Rate 95.550 % 4.550 5
Mem Pipes Busy 20.796 % 0.004 5
Warp Cycles Per Issued Instruction 28.766 cycle 0.038 5
Warp Cycles Per Executed Instruction 28.862 cycle 0.037 5
Avg. Active Threads Per Warp 30.460 0.000 5
Avg. Not Predicated Off Threads Per Warp 29.550 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 6.000 block 0.000 5
Block Limit Shared Mem 32.000 block 0.000 5
Block Limit Warps 8.000 block 0.000 5
Theoretical Active Warps per SM 48.000 warp 0.000 5
Theoretical Occupancy 75.000 % 0.000 5
Achieved Occupancy 61.950 % 0.068 5
Achieved Active Warps Per SM 39.648 warp 0.028 5
Analysis Rules
Rule Description
INF HighPipeUtilization ALU is the highest-utilized pipeline (20.3%) based on active cycles, taking into account the rates of its different instructions. It executes integer and logic operations. It is well-utilized, but should not be a bottleneck.
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 (75.0%) is limited by the number of required registers. The difference between calculated theoretical (75.0%) and measured achieved occupancy (62.4%) 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 367463.80 μs
Device Time 227.84 μs
Self CPU Time 73.08 μ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 367390.72 μs
Device Time 227.84 μs
Self CPU Time 119.48 μ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 366692.69 μs
Device Time 0.00 μs
Self CPU Time 131.92 μ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 364783.35 μs
Device Time 0.00 μs
Self CPU Time 364783.35 μ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
cudaLaunchKernel
CPU Time 667266.26 μs
Device Time 707.23 μs
Self CPU Time 667266.26 μs
Self Device Time 707.23 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
hybrid_optimized_kernel(float const*, float const*, float const*, float*, int, int, int)
CPU Time 0.00 μs
Device Time 246984.57 μs
Self CPU Time 0.00 μs
Self Device Time 246984.57 μ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 209336.48 μs
Device Time 596372.14 μs
Self CPU Time 12478.91 μ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 196859.86 μs
Device Time 596372.14 μs
Self CPU Time 16733.70 μs
Self Device Time 596372.14 μ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 596372.14 μs
Self CPU Time 0.00 μs
Self Device Time 596372.14 μ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
45315 warnings generated when compiling for host.
Suppressed 45347 warnings (45300 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/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:9:41 bugprone-easily-swappable-parameters
9 | __global__ void hybrid_optimized_kernel(const float* __restrict__ x,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~
10 | const float* __restrict__ weight,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
11 | const float* __restrict__ bias,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:9:67: note: the first parameter in the range is 'x'
9 | __global__ void hybrid_optimized_kernel(const float* __restrict__ x,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:11:68: note: the last parameter in the range is 'bias'
11 | const float* __restrict__ bias,
| ^~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:13:42: warning: 3 adjacent parameters of 'hybrid_optimized_kernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
13 | int batch_size,
| ^~~~~~~~~~~~~~~
14 | int in_features,
| ~~~~~~~~~~~~~~~~
15 | int out_features) {
| ~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:13:46: note: the first parameter in the range is 'batch_size'
13 | int batch_size,
| ^~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:15:46: note: the last parameter in the range is 'out_features'
15 | int out_features) {
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:16:21: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
16 | int batch_idx = blockIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:17:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
17 | int warp_id = threadIdx.x / WARP_SIZE;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:18:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
18 | int lane_id = threadIdx.x % WARP_SIZE;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:19:27: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
19 | int warps_per_block = blockDim.x / WARP_SIZE;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:20:26: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
20 | int base_out_group = blockIdx.y * (warps_per_block * TILE_SIZE);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:26:26: warning: result of multiplication in type 'int' is used as a pointer offset after an implicit widening conversion to type 'ptrdiff_t' [bugprone-implicit-widening-of-multiplication-result]
26 | const float* x_row = x + batch_idx * in_features;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:26:30: note: make conversion explicit to silence this warning
5 | const float* x_row = x + batch_idx * in_features;
| ^~~~~~~~~~~~~~~~~~~~~~~
| static_cast<ptrdiff_t>()
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:26:30: note: perform multiplication in a wider type
26 | const float* x_row = x + batch_idx * in_features;
| ^~~~~~~~~
| static_cast<ptrdiff_t>()
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:31:34: warning: result of multiplication in type 'int' is used as a pointer offset after an implicit widening conversion to type 'ptrdiff_t' [bugprone-implicit-widening-of-multiplication-result]
31 | const float* w_row = weight + current_out * in_features;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:31:43: note: make conversion explicit to silence this warning
31 | const float* w_row = weight + current_out * in_features;
| ^~~~~~~~~~~~~~~~~~~~~~~~~
| static_cast<ptrdiff_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:31:43: note: perform multiplication in a wider type
31 | const float* w_row = weight + current_out * in_features;
| ^~~~~~~~~~~
| static_cast<ptrdiff_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:67:54: warning: the parameter 'x' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
67 | torch::Tensor hybrid_optimized_forward(torch::Tensor x,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:68:54: warning: the parameter 'weight' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
68 | torch::Tensor weight,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:69:54: warning: the parameter 'bias' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
69 | torch::Tensor bias) {
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:74:22: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
74 | int batch_size = x.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:75:23: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
75 | int in_features = x.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s2_hybrid_optimized_kernel/base/base.cu:76:24: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
76 | int out_features = weight.size(0);
| ^