← Back to Leaderboard

The AI CUDA Engineer 👷

76_Gemm_Add_ReLUwarp_tile_hybrid_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 and tile sizes
#define WARP_SIZE 32
#define TILE_SIZE 4  // Each warp computes TILE_SIZE outputs simultaneously

// Hybrid kernel: combines benefits of multi-output tiling (minimizing kernel launches) from Kernel 1
// and the simplicity of single-warp dot product from Kernel 2. It uses vectorized loads via float4, __ldg() 
// for read-only caching, and warp-level shuffle reductions to eliminate explicit synchronization.

__global__ void warp_tile_hybrid_kernel(const float* __restrict__ x,
                                          const float* __restrict__ weight,
                                          const float* __restrict__ bias,
                                          float* __restrict__ out,
                                          int in_features,
                                          int out_features) {
    // Determine warp and lane indices
    int warp_id = threadIdx.x / WARP_SIZE;
    int lane_id = threadIdx.x % WARP_SIZE;

    // Each block contains several warps
    int warps_per_block = blockDim.x / WARP_SIZE;

    // Batch index is driven by blockIdx.x
    int batch_idx = blockIdx.x;
    
    // blockIdx.y covers groups of output features, each group has (warps_per_block * TILE_SIZE) outputs
    int group_idx = blockIdx.y;
    int base_out_group = group_idx * (warps_per_block * TILE_SIZE);

    // Each warp computes a tile of TILE_SIZE consecutive outputs
    int out_base = base_out_group + warp_id * TILE_SIZE;
    if (out_base >= out_features) return;

    // Initialize accumulators for each output computed by the warp tile
    float sums[TILE_SIZE];
#pragma unroll
    for (int t = 0; t < TILE_SIZE; t++) {
        sums[t] = 0.0f;
    }

    // Pointer to the current input row
    const float* x_row = x + batch_idx * in_features;

    // For each output in the tile, compute the dot product
#pragma unroll
    for (int t = 0; t < TILE_SIZE; t++) {
        int current_out = out_base + t;
        if (current_out < out_features) {
            // Get the corresponding weight row for current output feature
            const float* w_row = weight + current_out * in_features;

            // Use vectorized loads: float4 accesses (requires in_features to be a multiple of 4 for full efficiency)
            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);

            // Each thread in the warp handles a strided portion of the vectorized dot-product
            for (int i = lane_id; i < vec_count; i += WARP_SIZE) {
                float4 x_val = __ldg(x_vec + i);
                float4 w_val = __ldg(w_vec + i);
                sums[t] += x_val.x * w_val.x + x_val.y * w_val.y + x_val.z * w_val.z + x_val.w * w_val.w;
            }

            // Process remaining elements, if any, with scalar loads
            int start = vec_count * 4;
            for (int i = start + lane_id; i < in_features; i += WARP_SIZE) {
                sums[t] += __ldg(x_row + i) * __ldg(w_row + i);
            }
        }
    }

    // Perform warp-level reduction with shuffle instructions for each output in the tile
#pragma unroll
    for (int t = 0; t < TILE_SIZE; t++) {
        for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
            sums[t] += __shfl_down_sync(0xffffffff, sums[t], offset);
        }
    }

    // Lane 0 of each warp writes the final result, adding bias and applying ReLU activation
    if (lane_id == 0) {
        for (int t = 0; t < TILE_SIZE; t++) {
            int current_out = out_base + t;
            if (current_out < out_features) {
                float result = sums[t] + __ldg(bias + current_out);
                out[batch_idx * out_features + current_out] = (result > 0.0f) ? result : 0.0f;
            }
        }
    }
}

// Host function: prepares and launches the hybrid kernel
// Grid configuration: blockIdx.x spans the batch dimension and blockIdx.y spans groups of output features
// Each block is organized with a fixed number of threads such that each warp handles TILE_SIZE outputs

torch::Tensor warp_tile_hybrid_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());

    // Tunable parameter: Number of warps per block
    const int warps_per_block = 8;
    const int threads_per_block = warps_per_block * WARP_SIZE;

    // Calculate number of groups along the output feature dimension
    int groups_y = (out_features + (warps_per_block * TILE_SIZE) - 1) / (warps_per_block * TILE_SIZE);

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

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

    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &warp_tile_hybrid_forward, "Hybrid GEMM with bias and ReLU using warp-level tiling and optimized __ldg() loads (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 1.554 inst/cycle 0.000 5
Executed Ipc Elapsed 1.376 inst/cycle 0.000 5
Issue Slots Busy 38.948 % 0.043 5
Issued Ipc Active 1.556 inst/cycle 0.000 5
SM Busy 38.948 % 0.043 5
Memory Throughput 81756426925.122 byte/second 70824613530397584.000 5
Mem Busy 70.498 % 0.068 5
Max Bandwidth 68.924 % 0.065 5
L1/TEX Hit Rate 57.864 % 0.254 5
L2 Hit Rate 97.086 % 10.244 5
Mem Pipes Busy 23.138 % 0.007 5
Warp Cycles Per Issued Instruction 24.052 cycle 0.008 5
Warp Cycles Per Executed Instruction 24.146 cycle 0.008 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 5.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 40.000 warp 0.000 5
Theoretical Occupancy 62.500 % 0.000 5
Achieved Occupancy 58.550 % 0.015 5
Achieved Active Warps Per SM 37.468 warp 0.006 5
Analysis Rules
Rule Description
INF HighPipeUtilization ALU is the highest-utilized pipeline (22.6%) 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 (62.5%) is limited by the number of required registers. 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 412777.48 μs
Device Time 252.19 μs
Self CPU Time 59.10 μ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 412718.38 μs
Device Time 252.19 μs
Self CPU Time 136.45 μ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 411917.48 μs
Device Time 0.00 μs
Self CPU Time 136.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
cudaDeviceGetStreamPriorityRange
CPU Time 409805.99 μs
Device Time 0.00 μs
Self CPU Time 409805.99 μ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 398844.57 μs
Device Time 548.60 μs
Self CPU Time 398844.57 μs
Self Device Time 548.60 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
warp_tile_hybrid_kernel(float const*, float const*, float const*, float*, int, int)
CPU Time 0.00 μs
Device Time 168777.67 μs
Self CPU Time 0.00 μs
Self Device Time 168777.67 μ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 114485.17 μs
Device Time 475072.12 μs
Self CPU Time 16636.15 μ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 97850.78 μs
Device Time 475072.12 μs
Self CPU Time 22156.84 μs
Self Device Time 475072.12 μ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 475072.12 μs
Self CPU Time 0.00 μs
Self Device Time 475072.12 μ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_s1_warp_tile_hybrid/base/base.cu:14:41 bugprone-easily-swappable-parameters
14 | __global__ void warp_tile_hybrid_kernel(const float* __restrict__ x,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~
15 | const float* __restrict__ weight,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
16 | const float* __restrict__ bias,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:14:67: note: the first parameter in the range is 'x'
14 | __global__ void warp_tile_hybrid_kernel(const float* __restrict__ x,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:16:69: note: the last parameter in the range is 'bias'
16 | const float* __restrict__ bias,
| ^~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:18:43: warning: 2 adjacent parameters of 'warp_tile_hybrid_kernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
18 | int in_features,
| ^~~~~~~~~~~~~~~~
19 | int out_features) {
| ~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:18:47: note: the first parameter in the range is 'in_features'
18 | int in_features,
| ^~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:19:47: note: the last parameter in the range is 'out_features'
19 | int out_features) {
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:21:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
21 | 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_s1_warp_tile_hybrid/base/base.cu:22:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
22 | 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_s1_warp_tile_hybrid/base/base.cu:25:27: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
25 | 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_s1_warp_tile_hybrid/base/base.cu:28:21: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
28 | int batch_idx = blockIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:31:21: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
31 | int group_idx = blockIdx.y;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:46: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]
46 | 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_s1_warp_tile_hybrid/base/base.cu:46: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_s1_warp_tile_hybrid/base/base.cu:46:30: note: perform multiplication in a wider type
46 | 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_s1_warp_tile_hybrid/base/base.cu:54: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]
54 | 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_s1_warp_tile_hybrid/base/base.cu:54:43: note: make conversion explicit to silence this warning
54 | 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_s1_warp_tile_hybrid/base/base.cu:54:43: note: perform multiplication in a wider type
54 | 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_s1_warp_tile_hybrid/base/base.cu:100: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]
100 | torch::Tensor warp_tile_hybrid_forward(torch::Tensor x,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:101:56: 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]
101 | torch::Tensor weight,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:102:56: 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]
102 | torch::Tensor bias) {
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:107:22: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
107 | int batch_size = x.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:108:23: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
108 | int in_features = x.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s1_warp_tile_hybrid/base/base.cu:109:24: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
109 | int out_features = weight.size(0);
| ^