You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.

Key optimization techniques used in this implementation:
1.Operator Fusion: Fused normalization + affine transformation + sigmoid gating into a single kernel
2.Vectorized Memory Access: Utilizes float4 vector loads/stores for improved memory bandwidth utilization
3.Fast Math Functions: Employs optimized mathematical operations including rsqrtf, fmaf, and custom fast sigmoid
4.Read-Only Cache Optimization: Uses __ldg()intrinsic for constant memory access patterns
5.Loop Unrolling: Implements compile-time loop unrolling for reduced instruction overhead
6.Memory Access Coalescing: Organized thread-block mapping for optimal global memory access patterns
7.Exact Statistical Computation: Maintains numerical precision with proper variance calculation while optimizing performance

The custom CUDA implementation provides significant performance improvements over the native PyTorch version by eliminating intermediate tensor allocations and leveraging GPU-specific optimizations.
Specific Technical Optimizations:
Memory Hierarchy Optimization:
1.Global Memory: Vectorized loads/stores using float4for 4x bandwidth improvement
2.Constant Cache: __ldg()for parameter access (gamma, beta, v)
3.Register Utilization: Extensive use of registers for temporary variables

Computational Optimizations:
1.Fast Inverse Square Root: rsqrtf(variance + eps)instead of 1.0f/sqrtf()
2.Fused Multiply-Add: fmaf()instructions for affine transformation
3.Optimized Sigmoid: Custom fast_sigmoid()using __fdividefand __expf

Parallelism Strategy:
1.Grid Structure: One block per channel-instance combination (N×C blocks)
2.Block Configuration: 256 threads per block for optimal occupancy
3.Workload Balancing: Dynamic workload distribution across threads with unrolling

Numerical Precision:
1.Maintains mathematical equivalence with reference implementation
2.Proper handling of epsilon for numerical stability
3.Exact statistical computation preserved despite performance optimizations

The implementation demonstrates how custom CUDA kernels can dramatically accelerate normalization layers while maintaining full functional compatibility with standard PyTorch operations.
"""
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义维度常量
N, C, H, W = 32, 64, 56, 56
EPS = 1e-6


class EvoNormS0(nn.Module):
    """
    EvoNorm-S0: Evolving Normalization-Activation Layers (Sample-based, no batch dependency)

    公式:
    v = Var(x) = mean(x^2) - mean(x)^2
    y = x / sqrt(v + eps) * gamma + beta
    y = y * sigmoid(x * w)

    其中 gamma, beta, w 是可学习参数
    """

    def __init__(self, num_channels, eps, nonlinear=True):
        super().__init__()
        self.eps = eps
        self.nonlinear = nonlinear  # 是否使用非线性激活

        # 可学习的缩放和偏移参数（类似 BatchNorm）
        self.gamma = nn.Parameter(torch.ones(1, num_channels, 1, 1))
        self.beta = nn.Parameter(torch.zeros(1, num_channels, 1, 1))

        # 非线性门控参数
        if self.nonlinear:
            self.v = nn.Parameter(torch.ones(1, num_channels, 1, 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. 计算实例级方差
        # var = E[x^2] - E[x]^2
        x_sq_mean = torch.mean(x * x, dim=[2, 3], keepdim=True)
        x_mean = torch.mean(x, dim=[2, 3], keepdim=True)
        var = x_sq_mean - x_mean * x_mean

        # 2. 归一化
        x_normalized = x / torch.sqrt(var + self.eps)

        # 3. 仿射变换
        y = x_normalized * self.gamma + self.beta

        # 4. 非线性门控（可选）
        if self.nonlinear:
            y = y * torch.sigmoid(x * self.v)

        return y


class EvoNormB0(nn.Module):
    """
    EvoNorm-B0: Evolving Normalization-Activation Layers (Batch-based)

    公式:
    Instance Norm: x_in = (x - mean(x)) / sqrt(var(x) + eps)
    Batch Norm stats: rolling_var = momentum * rolling_var + (1-momentum) * batch_var
    y = x_in * gamma + beta
    y = y * sigmoid(x * w)
    """

    def __init__(self, num_channels, eps, momentum=0.1, nonlinear=True):
        super().__init__()
        self.eps = eps
        self.momentum = momentum
        self.nonlinear = nonlinear

        # 可学习参数
        self.gamma = nn.Parameter(torch.ones(1, num_channels, 1, 1))
        self.beta = nn.Parameter(torch.zeros(1, num_channels, 1, 1))

        # 非线性门控参数
        if self.nonlinear:
            self.v = nn.Parameter(torch.ones(1, num_channels, 1, 1))

        # 运行时统计量（用于推理）
        self.register_buffer('running_var', torch.ones(1, num_channels, 1, 1))
        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.training:
            # 训练模式：计算当前批次的统计量
            # 1. 实例归一化
            x_mean = torch.mean(x, dim=[2, 3], keepdim=True)
            x_var = torch.var(x, dim=[2, 3], keepdim=True, unbiased=False)

            # 2. 更新运行统计量（跨批次的方差）
            batch_var = torch.mean(x_var, dim=0, keepdim=True)
            with torch.no_grad():
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var
                self.num_batches_tracked += 1

            # 3. 归一化
            x_normalized = (x - x_mean) / torch.sqrt(x_var + self.eps)
        else:
            # 推理模式：使用运行统计量
            x_mean = torch.mean(x, dim=[2, 3], keepdim=True)
            x_normalized = (x - x_mean) / torch.sqrt(self.running_var + self.eps)

        # 4. 仿射变换
        y = x_normalized * self.gamma + self.beta

        # 5. 非线性门控
        if self.nonlinear:
            y = y * torch.sigmoid(x * self.v)

        return y


class Model(nn.Module):
    """
    EvoNorm 模型包装器
    默认使用 EvoNorm-S0（无批次依赖，更适合小批量）
    """

    def __init__(self, evonorm_gamma, evonorm_beta, evonorm_v=None, use_b0=False):
        super().__init__()

        # 选择 EvoNorm 变体
        if use_b0:
            self.evonorm = EvoNormB0(C, EPS, nonlinear=(evonorm_v is not None))
        else:
            self.evonorm = EvoNormS0(C, EPS, nonlinear=(evonorm_v is not None))

        # 初始化参数
        with torch.no_grad():
            self.evonorm.gamma.data.copy_(evonorm_gamma.view(1, C, 1, 1))
            self.evonorm.beta.data.copy_(evonorm_beta.view(1, C, 1, 1))

            if evonorm_v is not None and self.evonorm.nonlinear:
                self.evonorm.v.data.copy_(evonorm_v.view(1, C, 1, 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.evonorm(x)


def get_inputs():
    """生成测试输入"""
    x = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x]


def get_init_inputs():
    """
    生成初始化参数
    返回 [gamma, beta, v]
    """
    evonorm_gamma = torch.ones(1, C, 1, 1)
    evonorm_beta = torch.zeros(1, C, 1, 1)
    evonorm_v = torch.ones(1, C, 1, 1)  # 门控参数
    return [evonorm_gamma, evonorm_beta, evonorm_v]