← Back to Leaderboard

The AI CUDA Engineer 👷

36_RMSNorm_ • 36_rmsnorm_even_workload_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_even_workload_base
CUDA Speedup (Native) 2.808x
CUDA Speedup (Compile) 2.683x
CUDA Runtime 0.189 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>

// Define block dimensions for balanced workload distribution
#define OFFSETS_PER_BLOCK 32
#define THREADS_FEATURE 8

// Each column (threadIdx.x) in the 2D block processes one (batch, offset) pair.
// Threads in the column (indexed by threadIdx.y) cooperatively compute the sum of squares
// across the feature dimension and then normalize the corresponding elements.

template <typename scalar_t>
__global__ void rms_norm_even_workload_kernel(
    const scalar_t* __restrict__ input,
    scalar_t* __restrict__ output,
    const int total_offsets,  // batch_size * numel_per_batch
    const int num_features,
    const int numel_per_batch,
    const float eps
) {
    // Calculate the global offset index corresponding to a (batch, offset) pair
    int global_offset = blockIdx.x * OFFSETS_PER_BLOCK + threadIdx.x;
    if (global_offset >= total_offsets) return;

    // Determine the batch id and the offset within the batch
    int batch_id = global_offset / numel_per_batch;
    int offset = global_offset % numel_per_batch;
    int base = batch_id * num_features * numel_per_batch;

    // Shared memory for reduction: size = OFFSETS_PER_BLOCK * THREADS_FEATURE
    __shared__ scalar_t sdata[OFFSETS_PER_BLOCK * THREADS_FEATURE];

    // Each thread in the column computes a partial sum over a subset of feature indices
    scalar_t partial_sum = 0;
    for (int f = threadIdx.y; f < num_features; f += THREADS_FEATURE) {
        int pos = base + f * numel_per_batch + offset;
        scalar_t val = input[pos];
        partial_sum += val * val;
    }

    // Store the partial sum in shared memory. Shared memory is laid out as [THREADS_FEATURE][OFFSETS_PER_BLOCK]
    int smem_index = threadIdx.y * OFFSETS_PER_BLOCK + threadIdx.x;
    sdata[smem_index] = partial_sum;
    __syncthreads();

    // Perform reduction along the feature dimension (vertical reduction within the column)
    for (int stride = THREADS_FEATURE / 2; stride > 0; stride /= 2) {
        if (threadIdx.y < stride) {
            sdata[smem_index] += sdata[(threadIdx.y + stride) * OFFSETS_PER_BLOCK + threadIdx.x];
        }
        __syncthreads();
    }

    // Thread with threadIdx.y == 0 in each column now holds the complete sum of squares
    scalar_t rms;
    if (threadIdx.y == 0) {
        scalar_t sumsq = sdata[threadIdx.x];
        rms = sqrt(sumsq / num_features + eps);
        // Store the computed rms in shared memory for use by all threads in this column
        sdata[threadIdx.x] = rms;
    }
    __syncthreads();
    rms = sdata[threadIdx.x];

    // Normalization: each thread in the column normalizes a subset of feature elements
    for (int f = threadIdx.y; f < num_features; f += THREADS_FEATURE) {
        int pos = base + f * numel_per_batch + offset;
        scalar_t val = input[pos];
        output[pos] = val / rms;
    }
}

// CUDA forward function with a 2D block layout for even workload distribution
torch::Tensor rms_norm_cuda_forward_even_workload(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);

    int numel_per_batch = 1;
    for (int i = 2; i < input.dim(); i++) {
        numel_per_batch *= input.size(i);
    }

    // Total number of (batch, offset) pairs to process
    int total_offsets = batch_size * numel_per_batch;

    // Define block dimensions: each block has OFFSETS_PER_BLOCK columns and THREADS_FEATURE rows
    dim3 block(OFFSETS_PER_BLOCK, THREADS_FEATURE);
    int grid_x = (total_offsets + OFFSETS_PER_BLOCK - 1) / OFFSETS_PER_BLOCK;
    dim3 grid(grid_x);

    AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "rms_norm_cuda_even_workload", ([&] {
        rms_norm_even_workload_kernel<scalar_t><<<grid, block>>>(
            input.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            total_offsets,
            num_features,
            numel_per_batch,
            eps
        );
    }));

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &rms_norm_cuda_forward_even_workload, "RMS normalization forward with balanced workload (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 1.412 inst/cycle 0.000 5
Executed Ipc Elapsed 1.374 inst/cycle 0.000 5
Issue Slots Busy 35.302 % 0.054 5
Issued Ipc Active 1.412 inst/cycle 0.000 5
SM Busy 35.302 % 0.054 5
Memory Throughput 2890460772180.298 byte/second 74321090274980085760.000 5
Mem Busy 51.122 % 0.017 5
Max Bandwidth 86.242 % 0.064 5
L1/TEX Hit Rate 15.124 % 0.002 5
L2 Hit Rate 57.104 % 0.016 5
Mem Pipes Busy 27.344 % 0.004 5
Warp Cycles Per Issued Instruction 38.642 cycle 0.041 5
Warp Cycles Per Executed Instruction 38.654 cycle 0.042 5
Avg. Active Threads Per Warp 32.000 0.000 5
Avg. Not Predicated Off Threads Per Warp 27.650 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 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 85.734 % 0.002 5
Achieved Active Warps Per SM 54.870 warp 0.001 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 (85.7%) 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 340305.94 μs
Device Time 27089.25 μs
Self CPU Time 43.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 340262.84 μs
Device Time 27089.25 μs
Self CPU Time 102.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
cudaLaunchKernel
CPU Time 1160076.58 μs
Device Time 13631.91 μs
Self CPU Time 1160076.58 μs
Self Device Time 13631.91 μ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_even_workload_kernel<float>(float const*, float*, int, int, int, float)
CPU Time 0.00 μs
Device Time 935791.75 μs
Self CPU Time 0.00 μs
Self Device Time 935791.75 μ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 865449.64 μs
Device Time 393727.27 μs
Self CPU Time 8117.84 μ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 857333.53 μs
Device Time 393727.27 μs
Self CPU Time 12445.50 μs
Self Device Time 393727.27 μ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 393727.27 μs
Self CPU Time 0.00 μs
Self Device Time 393727.27 μ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 
45288 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/b9_s2_36_rmsnorm_even_workload/base/base.cu:18:5 bugprone-easily-swappable-parameters
18 | const int total_offsets, // batch_size * numel_per_batch
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
19 | const int num_features,
| ~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:18:15: note: the first parameter in the range is 'total_offsets'
18 | const int total_offsets, // batch_size * numel_per_batch
| ^~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:19:15: note: the last parameter in the range is 'num_features'
19 | const int num_features,
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:20:5: warning: 2 adjacent parameters of 'rms_norm_even_workload_kernel' of convertible types are easily swapped by mistake [bugprone-easily-swappable-parameters]
20 | const int numel_per_batch,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~
21 | const float eps
| ~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:20:15: note: the first parameter in the range is 'numel_per_batch'
20 | const int numel_per_batch,
| ^~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:21:17: note: the last parameter in the range is 'eps'
21 | const float eps
| ^~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:21: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')
21 | const float eps
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:24:25: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
24 | int global_offset = blockIdx.x * OFFSETS_PER_BLOCK + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:37:18: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
37 | for (int f = threadIdx.y; f < num_features; f += THREADS_FEATURE) {
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:44:22: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
44 | int smem_index = threadIdx.y * OFFSETS_PER_BLOCK + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:68:18: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
68 | for (int f = threadIdx.y; f < num_features; f += THREADS_FEATURE) {
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:79:28: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
79 | const int batch_size = input.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:80:30: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
80 | const int num_features = input.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:84:28: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
84 | numel_per_batch *= input.size(i);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_36/b9_s2_36_rmsnorm_even_workload/base/base.cu:95: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]
95 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "rms_norm_cuda_even_workload", ([&] {
| ^
/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__, \
| ^