← Back to Leaderboard

The AI CUDA Engineer 👷

36_RMSNorm_ • 36_rmsnorm_unroll_opt_base

Level 1 • Task 36
import torch
import torch.nn as nn
import torch.nn.functional as F


def module_fn(x: torch.Tensor, eps: float) -> torch.Tensor:
    """
    Applies RMS Normalization to the input tensor.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, num_features, *)
        eps (float): Small value added to denominator for numerical stability

    Returns:
        torch.Tensor: Output tensor with RMS Normalization applied
    """
    rms = torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + eps)
    return x / rms


class Model(nn.Module):
    """
    Simple model that performs RMS Normalization.
    """

    def __init__(self, num_features: int, eps: float):
        """
        Initializes the RMSNorm layer.

        Args:
            num_features (int): Number of features in the input tensor
            eps (float): Small value added to denominator for numerical stability
        """
        super(Model, self).__init__()
        self.eps = eps

    def forward(self, x: torch.Tensor, fn=module_fn) -> torch.Tensor:
        """
        Forward pass that calls module_fn.

        Args:
            x (torch.Tensor): Input tensor
            fn: Function to call, defaults to module_fn

        Returns:
            torch.Tensor: Output of module_fn
        """
        return fn(x, self.eps)


batch_size = 16
features = 64
dim1 = 256
dim2 = 256
eps = 1e-5


def get_inputs():
    x = torch.randn(batch_size, features, dim1, dim2)
    return [x]


def get_init_inputs():
    return [features, eps]

import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Simple model that performs RMS Normalization.
    """

    def __init__(self, num_features: int, eps: float = 1e-5):
        """
        Initializes the RMSNorm layer.

        Args:
            num_features (int): Number of features in the input tensor.
            eps (float, optional): A small value added to the denominator to avoid division by zero. Defaults to 1e-5.
        """
        super(Model, self).__init__()
        self.num_features = num_features
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies RMS Normalization to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, num_features, *).

        Returns:
            torch.Tensor: Output tensor with RMS Normalization applied, same shape as input.
        """
        # Calculate the RMS along the feature dimension
        rms = torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + self.eps)

        # Normalize the input by dividing by the RMS
        return x / rms


batch_size = 16
features = 64
dim1 = 256
dim2 = 256
eps = 1e-5


def get_inputs():
    x = torch.randn(batch_size, features, dim1, dim2)
    return [x]


def get_init_inputs():
    return [features, eps]
Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name 36_RMSNorm_
Level ID 1
Task ID 36
Kernel Name 36_rmsnorm_unroll_opt_base
CUDA Speedup (Native) 1.937x
CUDA Speedup (Compile) 1.851x
CUDA Runtime 0.274 ms
PyTorch Runtime (Native) 0.531 ms
PyTorch Runtime (Compile) 0.507 ms
Correct True
Max Diff (vs. Reference) 0.000000
Model o3-mini-2025-01-31
Temperature 1.00
View Experiment Progress Details

Related Kernels (Level 1, Task 36 • 36_RMSNorm_)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 36_rmsnorm_even_workload_base 0.19 2.81 2.68
🥇 36_rmsnorm_optimized_indexing_base_base 0.19 2.81 2.68
🥉 36_rmsnorm_optimized_indexing_base_base 0.19 2.79 2.67
4 modular_rms_norm_edit_1 0.27 1.95 1.86
4 modular_rms_norm_base 0.27 1.95 1.86
4 combined_grid_unroll_base 0.27 1.95 1.86
7 variable_block_size_rms_norm_edit_1 0.27 1.94 1.85
7 36_rmsnorm_ldg_aligned_opt_base 0.27 1.94 1.85
7 36_rmsnorm_modular_funcs_base 0.27 1.94 1.85
7 36_RMSNorm_uniform_control_base_base 0.27 1.94 1.85
7 balanced_workload_rms_base 0.27 1.94 1.85
7 36_rmsnorm_unroll_opt_base 0.27 1.94 1.85
7 unrolled_rms_norm_edit_1 0.27 1.94 1.85
7 unrolled_rms_norm_base 0.27 1.94 1.85
7 variable_block_size_rms_norm_base 0.27 1.94 1.85
7 36_rmsnorm_modular_inlined_base 0.27 1.94 1.85
7 36_RMSNorm_optimized_block_size_base 0.27 1.94 1.85
18 36_RMSNorm_stride_loop_base 0.28 1.93 1.84
18 efficient_rms_norm_kernel_base 0.28 1.93 1.84
18 36_RMSNorm_no_divergence_base 0.28 1.93 1.84 
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// Kernel with manual loop unrolling using #pragma unroll

template <typename scalar_t>
__global__ void rms_norm_kernel_unroll(
    const scalar_t* __restrict__ input,
    scalar_t* __restrict__ output,
    const int total_work,    // batch_size * numel_per_batch
    const int num_features,
    const int numel_per_batch,
    const float eps
) {
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int stride = blockDim.x * gridDim.x;

    // Each thread processes multiple work-items in strides
    for (int index = tid; index < total_work; index += stride) {
        const int batch_id = index / numel_per_batch;
        const int offset = index % numel_per_batch;
        const int batch_offset = batch_id * num_features * numel_per_batch;

        scalar_t sumsq = static_cast<scalar_t>(0);

        // Unroll the loop for summing the squares
        #pragma unroll
        for (int feat = 0; feat < num_features; feat++) {
            int idx = batch_offset + feat * numel_per_batch + offset;
            scalar_t val = input[idx];
            sumsq += val * val;
        }

        // Compute the RMS value with epsilon for stability
        scalar_t rms = sqrt(sumsq / num_features + eps);

        // Unroll the loop for normalizing the values
        #pragma unroll
        for (int feat = 0; feat < num_features; feat++) {
            int idx = batch_offset + feat * numel_per_batch + offset;
            output[idx] = input[idx] / rms;
        }
    }
}

// CUDA forward function

torch::Tensor rms_norm_cuda_forward(torch::Tensor input, float eps) {
    auto output = torch::empty_like(input);

    const int batch_size = input.size(0);
    const int num_features = input.size(1);

    // Calculate number of elements per batch for dimensions beyond the first two
    int numel_per_batch = 1;
    for (int i = 2; i < input.dim(); i++) {
        numel_per_batch *= input.size(i);
    }

    int total_work = batch_size * numel_per_batch;
    const int threads_per_block = 256;
    const int blocks = (total_work + threads_per_block - 1) / threads_per_block;

    AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "rms_norm_cuda_unroll", ([&] {
        rms_norm_kernel_unroll<scalar_t><<<blocks, threads_per_block>>>(
            input.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            total_work,
            num_features,
            numel_per_batch,
            eps
        );
    }));

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &rms_norm_cuda_forward, "RMS normalization forward with loop unrolling (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 0.440 inst/cycle 0.000 5
Executed Ipc Elapsed 0.430 inst/cycle 0.000 5
Issue Slots Busy 11.076 % 0.001 5
Issued Ipc Active 0.440 inst/cycle 0.000 5
SM Busy 12.578 % 0.001 5
Memory Throughput 2960479460631.492 byte/second 33897395436009021440.000 5
Mem Busy 48.648 % 0.009 5
Max Bandwidth 88.324 % 0.030 5
L1/TEX Hit Rate 0.302 % 0.000 5
L2 Hit Rate 34.248 % 0.002 5
Mem Pipes Busy 11.408 % 0.000 5
Warp Cycles Per Issued Instruction 134.876 cycle 0.372 5
Warp Cycles Per Executed Instruction 135.050 cycle 0.375 5
Avg. Active Threads Per Warp 32.000 0.000 5
Avg. Not Predicated Off Threads Per Warp 31.470 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 8.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 64.000 warp 0.000 5
Theoretical Occupancy 100.000 % 0.000 5
Achieved Occupancy 93.624 % 0.001 5
Achieved Active Warps Per SM 59.920 warp 0.000 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.
INF Occupancy This kernel's theoretical occupancy is not impacted by any block limit.
Operation / Metric Value Unit
aten::to
CPU Time 534114.70 μs
Device Time 27494.04 μs
Self CPU Time 33.93 μ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 534080.77 μs
Device Time 27494.04 μs
Self CPU Time 79.41 μ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 1484113.73 μs
Device Time 12990.18 μs
Self CPU Time 1484113.73 μs
Self Device Time 12990.18 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
void rms_norm_kernel_unroll<float>(float const*, float*, int, int, int, float)
CPU Time 0.00 μs
Device Time 1301550.89 μs
Self CPU Time 0.00 μs
Self Device Time 1301550.89 μ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 1204027.80 μs
Device Time 374963.05 μs
Self CPU Time 7830.55 μ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 1196198.80 μs
Device Time 374963.05 μs
Self CPU Time 10916.06 μs
Self Device Time 374963.05 μ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 374963.05 μs
Self CPU Time 0.00 μs
Self Device Time 374963.05 μ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 
45286 warnings generated when compiling for host.
Suppressed 45322 warnings (45275 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/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:12:5 bugprone-easily-swappable-parameters
12 | const int total_work, // batch_size * numel_per_batch
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
13 | const int num_features,
| ~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:12:15: note: the first parameter in the range is 'total_work'
12 | const int total_work, // batch_size * numel_per_batch
| ^~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:13:15: note: the last parameter in the range is 'num_features'
13 | const int num_features,
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:14:5: warning: 2 adjacent parameters of 'rms_norm_kernel_unroll' of convertible types are easily swapped by mistake [bugprone-easily-swappable-parameters]
14 | const int numel_per_batch,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~
15 | const float eps
| ~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:14:15: note: the first parameter in the range is 'numel_per_batch'
14 | const int numel_per_batch,
| ^~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:15:17: note: the last parameter in the range is 'eps'
15 | const float eps
| ^~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:15:5: note: 'const int' and 'const float' may be implicitly converted: 'const int' (as 'int') -> 'const float' (as 'float'), 'const float' (as 'float') -> 'const int' (as 'int')
15 | const float eps
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:17:21: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
17 | const int tid = blockIdx.x * blockDim.x + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:18:24: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
18 | const int stride = blockDim.x * gridDim.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:53:28: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
53 | const int batch_size = input.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:54:30: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
54 | const int num_features = input.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:59:28: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
59 | numel_per_batch *= input.size(i);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b6_s2_36_rmsnorm_unroll_opt/base/base.cu:66: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]
66 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "rms_norm_cuda_unroll", ([&] {
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:246:19: note: expanded from macro 'AT_DISPATCH_FLOATING_TYPES_AND_HALF'
246 | TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES_AND_HALF(__VA_ARGS__))
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:240:3: note: expanded from macro 'AT_DISPATCH_CASE_FLOATING_TYPES_AND_HALF'
240 | 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__, \
| ^