Kernel Details - fused_sigmoid_groupnorm

import torch
import torch.nn as nn
import torch.nn.functional as F


def module_fn(
    x: torch.Tensor,
    conv_weight: torch.Tensor,
    conv_bias: torch.Tensor,
    bias: torch.Tensor,
    scale: torch.Tensor,
    group_norm_weight: torch.Tensor,
    group_norm_bias: torch.Tensor,
    num_groups: int,
) -> torch.Tensor:
    """
    Applies convolution, bias addition, scaling, sigmoid activation and group normalization.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width)
        conv_weight (torch.Tensor): Convolution weight tensor
        conv_bias (torch.Tensor): Convolution bias tensor
        bias (torch.Tensor): Bias tensor for addition
        scale (torch.Tensor): Scale tensor for multiplication
        group_norm_weight (torch.Tensor): Group norm weight tensor
        group_norm_bias (torch.Tensor): Group norm bias tensor
        num_groups (int): Number of groups for group normalization

    Returns:
        torch.Tensor: Output tensor after applying convolution, bias, scale, sigmoid and group norm
    """
    x = F.conv2d(x, conv_weight, bias=conv_bias)
    x = x + bias
    x = x * scale
    x = torch.sigmoid(x)
    x = F.group_norm(x, num_groups, group_norm_weight, group_norm_bias)
    return x


class Model(nn.Module):
    """
    Model that performs a convolution, adds a bias term, scales, applies sigmoid, and performs group normalization.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        num_groups,
        bias_shape,
        scale_shape,
    ):
        super(Model, self).__init__()
        conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.conv_weight = conv.weight
        self.conv_bias = nn.Parameter(
            conv.bias + torch.ones_like(conv.bias) * 0.02
        )  # make sure its nonzero
        self.bias = nn.Parameter(torch.randn(bias_shape) * 0.02)
        self.scale = nn.Parameter(torch.randn(scale_shape) * 0.02)
        group_norm = nn.GroupNorm(num_groups, out_channels)
        self.group_norm_weight = group_norm.weight
        self.group_norm_bias = nn.Parameter(
            group_norm.bias + torch.ones_like(group_norm.bias) * 0.02
        )  # make sure its nonzero

    def forward(self, x, num_groups, fn=module_fn):
        return fn(
            x,
            self.conv_weight,
            self.conv_bias,
            self.bias,
            self.scale,
            self.group_norm_weight,
            self.group_norm_bias,
            num_groups,
        )


batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
num_groups = 8
bias_shape = (out_channels, 1, 1)
scale_shape = (out_channels, 1, 1)


def get_inputs():
    return [torch.randn(batch_size, in_channels, height, width), num_groups]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape]

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, adds a bias term, scales, applies sigmoid, and performs group normalization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.conv.bias = nn.Parameter(self.conv.bias + torch.ones_like(self.conv.bias) * 0.02)
        self.bias = nn.Parameter(torch.randn(bias_shape)*0.02) 
        self.scale = nn.Parameter(torch.randn(scale_shape)*0.02)
        self.group_norm = nn.GroupNorm(num_groups, out_channels)
        self.group_norm.bias = nn.Parameter(self.group_norm.bias + torch.ones_like(self.group_norm.bias) * 0.02)

    def forward(self, x):
        x = self.conv(x)
        x = x + self.bias
        x = x * self.scale
        x = torch.sigmoid(x)
        x = self.group_norm(x)
        return x

batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
num_groups = 8
bias_shape = (out_channels, 1, 1)
scale_shape = (out_channels, 1, 1)

def get_inputs():
    return [torch.randn(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape]

Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name	21_Conv2d_Add_Scale_Sigmoid_GroupNorm
Level ID	2
Task ID	21
Kernel Name	fused_sigmoid_groupnorm_base
CUDA Speedup (Native)	1.630x
CUDA Speedup (Compile)	1.396x
CUDA Runtime	0.046 ms
PyTorch Runtime (Native)	0.075 ms
PyTorch Runtime (Compile)	0.064 ms
Correct	True
Max Diff (vs. Reference)	0.007000
Model	o3-mini-2025-01-31
Temperature	1.00

View Experiment Progress Details

Related Kernels (Level 2, Task 21 • 21_Conv2d_Add_Scale_Sigmoid_GroupNorm)

Rank	Kernel Name	Runtime (ms)	Speedup Native	Speedup Compile
🥇	shared_memory_coalesced_access_kernel_base	0.04	1.79	1.53
🥈	fused_elem_groupnorm_reduced_sync_base_base	0.04	1.74	1.49
🥈	fused_optimized_warp_base	0.04	1.74	1.49
🥈	optimized_memory_access_kernel_base	0.04	1.74	1.49
🥈	atomic_optimized_kernel_base_base	0.04	1.74	1.49
🥈	optimized_memory_access_kernel_base	0.04	1.74	1.49
🥈	optimized_fused_kernel_base	0.04	1.74	1.49
8	fused_warp_reduce_groupnorm_base	0.04	1.70	1.46
8	fused_warp_groupnorm_base	0.04	1.70	1.46
8	shared_memory_reuse_kernel_base	0.04	1.70	1.46
8	fused_lockfree_groupnorm_base_base	0.04	1.70	1.46
8	fused_stride_kernel_base	0.04	1.70	1.46
8	fused_elem_groupnorm_no_atomic_base	0.04	1.70	1.46
8	fused_elem_groupnorm_min_sync_base_base	0.04	1.70	1.46
8	unrolled_fused_kernel_base_base	0.04	1.70	1.46
16	optimized_modular_kernel_base	0.04	1.67	1.43
16	fused_strided_groupnorm_base_base	0.04	1.67	1.43
18	fused_sigmoid_groupnorm_base	0.05	1.63	1.40
19	fused_sigmoid_groupnorm_base	0.05	1.60	1.37
19	block_size_optimized_kernel_base_base	0.05	1.60	1.37

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor.")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous.")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

// Fused kernel that computes sigmoid activation and applies group normalization
// in two passes: first, compute activation and reduce for mean/variance; then, normalize
__global__ void fused_sigmoid_groupnorm_kernel(
    const float* __restrict__ x,
    float* __restrict__ y,
    const float* __restrict__ bias,
    const float* __restrict__ scale,
    const float* __restrict__ gn_weight,
    const float* __restrict__ gn_bias,
    int N, int C, int H, int W,
    int num_groups,
    bool bias_broadcast,
    bool scale_broadcast,
    float eps) {

    // Allocate shared memory for reduction: one array for the sum, one for sum of squares
    extern __shared__ float shared_mem[];
    float* shared_sum = shared_mem;              // size: blockDim.x
    float* shared_sum_sq = &shared_mem[blockDim.x]; // size: blockDim.x

    // Each block handles one group for one sample
    int block_id = blockIdx.x;
    int sample_idx = block_id / num_groups;
    int group_idx = block_id % num_groups;

    // Compute group parameters
    int channels_per_group = C / num_groups;
    int group_size = channels_per_group * H * W;
    int sample_offset = sample_idx * C * H * W;
    int group_channel_offset = group_idx * channels_per_group;

    int tid = threadIdx.x;
    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;

    // First pass: Compute sigmoid activated value and accumulate for group statistics
    int stride = blockDim.x;
    for (int i = tid; i < group_size; i += stride) {
        int c_local = i / (H * W);
        int hw = i % (H * W);
        int c = group_channel_offset + c_local;
        int index = sample_offset + c * H * W + hw;

        // Load input
        float x_val = x[index];
        // Select bias and scale depending on broadcast
        float bias_val = bias_broadcast ? bias[0] : bias[c];
        float scale_val = scale_broadcast ? scale[0] : scale[c];

        // Compute the transformed value and apply sigmoid activation
        float z = (x_val + bias_val) * scale_val;
        float activated = 1.0f / (1.0f + expf(-z));
        
        // Store the intermediate activated value
        y[index] = activated;
        
        local_sum += activated;
        local_sum_sq += activated * activated;
    }

    // Write local sums to shared memory
    shared_sum[tid] = local_sum;
    shared_sum_sq[tid] = local_sum_sq;
    __syncthreads();

    // Perform reduction in shared memory to compute total sum and sum of squares
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared_sum[tid] += shared_sum[tid + s];
            shared_sum_sq[tid] += shared_sum_sq[tid + s];
        }
        __syncthreads();
    }

    // Compute mean and variance for this group
    float mean = shared_sum[0] / group_size;
    float variance = (shared_sum_sq[0] / group_size) - (mean * mean);
    float inv_std = rsqrtf(variance + eps);
    __syncthreads();

    // Second pass: Normalize the activated values and apply gamma and beta
    for (int i = tid; i < group_size; i += stride) {
        int c_local = i / (H * W);
        int hw = i % (H * W);
        int c = group_channel_offset + c_local;
        int index = sample_offset + c * H * W + hw;

        float activated = y[index];
        float normalized = (activated - mean) * inv_std;
        float gamma = gn_weight[c];
        float beta = gn_bias[c];
        y[index] = gamma * normalized + beta;
    }
}

// Launcher for the fused kernel
void fused_sigmoid_groupnorm_cuda(
    at::Tensor x,
    at::Tensor bias,
    at::Tensor scale,
    at::Tensor y,
    at::Tensor gn_weight,
    at::Tensor gn_bias,
    int64_t num_groups,
    float eps) {

    const int N = x.size(0);
    const int C = x.size(1);
    const int H = x.size(2);
    const int W = x.size(3);

    // Configure kernel launch parameters
    int threads = 512; // Tunable depending on GPU
    int blocks = N * num_groups; // Each block handles one group's reduction
    int shared_mem_size = 2 * threads * sizeof(float);

    bool bias_broadcast = (bias.numel() == 1);
    bool scale_broadcast = (scale.numel() == 1);

    fused_sigmoid_groupnorm_kernel<<<blocks, threads, shared_mem_size>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        bias.data_ptr<float>(),
        scale.data_ptr<float>(),
        gn_weight.data_ptr<float>(),
        gn_bias.data_ptr<float>(),
        N, C, H, W, num_groups,
        bias_broadcast, scale_broadcast, eps);

    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "CUDA kernel failed: ", cudaGetErrorString(err));
}

// The forward function: performs convolution, then fused sigmoid activation and group normalization
at::Tensor module_fn_forward(
    at::Tensor x,
    at::Tensor conv_weight,
    at::Tensor conv_bias,
    at::Tensor bias,
    at::Tensor scale,
    at::Tensor gn_weight,
    at::Tensor gn_bias,
    int64_t num_groups) {

    CHECK_INPUT(x);
    CHECK_INPUT(conv_weight);
    if (conv_bias.defined()) CHECK_INPUT(conv_bias);
    CHECK_INPUT(bias);
    CHECK_INPUT(scale);
    CHECK_INPUT(gn_weight);
    CHECK_INPUT(gn_bias);

    // Perform convolution
    x = at::conv2d(x, conv_weight, conv_bias);

    // Allocate output tensor
    at::Tensor y = at::empty_like(x);

    // Launch the fused kernel
    fused_sigmoid_groupnorm_cuda(x, bias, scale, y, gn_weight, gn_bias, num_groups, 1e-5);

    return y;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &module_fn_forward, "Fused sigmoid activation and group normalization forward (CUDA)");
}

Performance Metrics

Metric	Value	Unit	Variance	Samples
Executed Ipc Active	2.482	inst/cycle	0.000	5
Executed Ipc Elapsed	2.102	inst/cycle	0.000	5
Issue Slots Busy	62.224	%	0.001	5
Issued Ipc Active	2.490	inst/cycle	0.000	5
SM Busy	62.224	%	0.001	5
Memory Throughput	370575348408.698	byte/second	6738898555083988992.000	5
Mem Busy	18.610	%	0.026	5
Max Bandwidth	21.718	%	0.024	5
L1/TEX Hit Rate	65.978	%	0.001	5
L2 Hit Rate	70.426	%	0.032	5
Mem Pipes Busy	24.326	%	0.030	5
Warp Cycles Per Issued Instruction	21.406	cycle	0.004	5
Warp Cycles Per Executed Instruction	21.456	cycle	0.004	5
Avg. Active Threads Per Warp	30.680		0.000	5
Avg. Not Predicated Off Threads Per Warp	27.620		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	4.000	block	0.000	5
Block Limit Shared Mem	12.000	block	0.000	5
Block Limit Warps	4.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	83.282	%	0.010	5
Achieved Active Warps Per SM	53.300	warp	0.005	5

Analysis Rules

Rule	Description
INF HighPipeUtilization	ALU is the highest-utilized pipeline (50.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.
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 (83.2%) 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.

Rule

Description

INF HighPipeUtilization

ALU is the highest-utilized pipeline (50.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.

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 (83.2%) 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::fill_
CPU Time	208876.13	μs
Device Time	1089399.18	μs
Self CPU Time	33858.99	μs
Self Device Time	1089399.18	μ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	234534.10	μs
Device Time	1089399.18	μs
Self CPU Time	25680.78	μ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::conv2d
CPU Time	1200607.33	μs
Device Time	389834.17	μs
Self CPU Time	24015.69	μ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::convolution
CPU Time	1176591.64	μs
Device Time	389834.17	μs
Self CPU Time	30257.71	μ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::_convolution
CPU Time	1146333.93	μs
Device Time	389834.17	μs
Self CPU Time	59996.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::cudnn_convolution
CPU Time	952930.14	μs
Device Time	277106.11	μs
Self CPU Time	219111.61	μs
Self Device Time	277106.11	μ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	1012474.11	μs
Device Time	28231.44	μs
Self CPU Time	1012474.11	μs
Self Device Time	28231.44	μ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	1089399.18	μs
Self CPU Time	0.00	μs
Self Device Time	1089399.18	μ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

45311 warnings generated when compiling for host.
Suppressed 45328 warnings (45281 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_21/b4_s3_fused_sigmoid_groupnorm/base/base.cu:7:35 bugprone-macro-parentheses

7 | #define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor.")

| ^

| ()

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b4_s3_fused_sigmoid_groupnorm/base/base.cu:8:41: warning: macro argument should be enclosed in parentheses [bugprone-macro-parentheses]

8 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous.")

| ^

| ()

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b4_s3_fused_sigmoid_groupnorm/base/base.cu:16:5: warning: 4 adjacent parameters of 'fused_sigmoid_groupnorm_kernel' of similar type ('const float *__restrict') are easily swapped by mistake [bugprone-easily-swappable-parameters]

16 | const float* __restrict__ bias,

| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

17 | const float* __restrict__ scale,

| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

18 | const float* __restrict__ gn_weight,

| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

19 | const float* __restrict__ gn_bias,

| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b4_s3_fused_sigmoid_groupnorm/base/base.cu:16:31: note: the first 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_21/b4_s3_fused_sigmoid_groupnorm/base/base.cu:19:31: note: the last parameter in the range is 'gn_bias'

19 | const float* __restrict__ gn_bias,

| ^~~~~~~

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b4_s3_fused_sigmoid_groupnorm/base/base.cu:20:5: warning: 2 adjacent parameters of 'fused_sigmoid_groupnorm_kernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]