```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# Let cuDNN pick the fastest convolution algorithm for the fixed shapes.
torch.backends.cudnn.benchmark = True

# Custom CUDA kernels:
#   1) group_stats_kernel: computes per-(sample, group) mean / rstd for GroupNorm.
#   2) fused_norm_act_lse_kernel / fused_norm_act_lse_vec4_kernel:
#      fuses GroupNorm affine -> Tanh -> HardSwish -> residual add -> channel-wise
#      LogSumExp (online, numerically stable) into a single pass over the conv output.
fused_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cfloat>

#define BLOCK_SIZE 256

__global__ void group_stats_kernel(const float* __restrict__ x,
                                   float* __restrict__ mean,
                                   float* __restrict__ rstd,
                                   int C, int G, int HW, int cpg, float eps) {
    int ng = blockIdx.x;
    int n = ng / G;
    int g = ng % G;
    long base = ((long)n * C + (long)g * cpg) * (long)HW;

    double s = 0.0;
    double sq = 0.0;
    for (int ci = 0; ci < cpg; ++ci) {
        const float* p = x + base + (long)ci * HW;
        for (int i = threadIdx.x; i < HW; i += blockDim.x) {
            float v = p[i];
            s += (double)v;
            sq += (double)v * (double)v;
        }
    }

    __shared__ double ssum[BLOCK_SIZE];
    __shared__ double ssq[BLOCK_SIZE];
    ssum[threadIdx.x] = s;
    ssq[threadIdx.x] = sq;
    __syncthreads();
    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            ssum[threadIdx.x] += ssum[threadIdx.x + stride];
            ssq[threadIdx.x] += ssq[threadIdx.x + stride];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        double cnt = (double)cpg * (double)HW;
        double m = ssum[0] / cnt;
        double var = ssq[0] / cnt - m * m;
        if (var < 0.0) var = 0.0;
        mean[ng] = (float)m;
        rstd[ng] = (float)(1.0 / sqrt(var + (double)eps));
    }
}

__device__ __forceinline__ float apply_chain(float v, float m, float r, float w, float b) {
    // GroupNorm affine
    float y = (v - m) * r * w + b;
    // Tanh
    float t = tanhf(y);
    // HardSwish: t * relu6(t + 3) / 6
    float hs = t * fminf(fmaxf(t + 3.0f, 0.0f), 6.0f) / 6.0f;
    // Residual add
    return v + hs;
}

__device__ __forceinline__ void lse_update(float r, float &m, float &s) {
    if (r > m) {
        s = s * expf(m - r) + 1.0f;
        m = r;
    } else {
        s += expf(r - m);
    }
}

__global__ void fused_norm_act_lse_kernel(const float* __restrict__ x,
                                          const float* __restrict__ mean,
                                          const float* __restrict__ rstd,
                                          const float* __restrict__ weight,
                                          const float* __restrict__ bias,
                                          float* __restrict__ out,
                                          long total, int C, int G, int HW, int cpg) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int n = (int)(idx / HW);
    int o = (int)(idx - (long)n * HW);

    const float* xp = x + ((long)n * C) * HW + o;
    const float* mp = mean + (long)n * G;
    const float* rp = rstd + (long)n * G;

    float m = -FLT_MAX;
    float s = 0.0f;
    for (int g = 0; g < G; ++g) {
        float mg = mp[g];
        float rg = rp[g];
        int c0 = g * cpg;
        for (int ci = 0; ci < cpg; ++ci) {
            int c = c0 + ci;
            float v = xp[(long)c * HW];
            float r = apply_chain(v, mg, rg, weight[c], bias[c]);
            lse_update(r, m, s);
        }
    }
    out[idx] = m + logf(s);
}

__global__ void fused_norm_act_lse_vec4_kernel(const float4* __restrict__ x,
                                               const float* __restrict__ mean,
                                               const float* __restrict__ rstd,
                                               const float* __restrict__ weight,
                                               const float* __restrict__ bias,
                                               float4* __restrict__ out,
                                               long total4, int C, int G, int HW4, int cpg) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total4) return;
    int n = (int)(idx / HW4);
    int o = (int)(idx - (long)n * HW4);

    const float4* xp = x + ((long)n * C) * HW4 + o;
    const float* mp = mean + (long)n * G;
    const float* rp = rstd + (long)n * G;

    float m0 = -FLT_MAX, m1 = -FLT_MAX, m2 = -FLT_MAX, m3 = -FLT_MAX;
    float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
    for (int g = 0; g < G; ++g) {
        float mg = mp[g];
        float rg = rp[g];
        int c0 = g * cpg;
        for (int ci = 0; ci < cpg; ++ci) {
            int c = c0 + ci;
            float4 v = xp[(long)c * HW4];
            float w = weight[c];
            float b = bias[c];
            lse_update(apply_chain(v.x, mg, rg, w, b), m0, s0);
            lse_update(apply_chain(v.y, mg, rg, w, b), m1, s1);
            lse_update(apply_chain(v.z, mg, rg, w, b), m2, s2);
            lse_update(apply_chain(v.w, mg, rg, w, b), m3, s3);
        }
    }
    out[idx] = make_float4(m0 + logf(s0), m1 + logf(s1),
                           m2 + logf(s2), m3 + logf(s3));
}

torch::Tensor fused_gn_act_lse(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                               int64_t groups, double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto x_c = x.contiguous();
    auto w_c = weight.contiguous();
    auto b_c = bias.contiguous();

    int N = x_c.size(0);
    int C = x_c.size(1);
    int H = x_c.size(2);
    int W = x_c.size(3);
    int HW = H * W;
    int G = (int)groups;
    TORCH_CHECK(C % G == 0, "channels must be divisible by groups");
    int cpg = C / G;

    auto opts = x_c.options();
    auto mean = torch::empty({N, G}, opts);
    auto rstd = torch::empty({N, G}, opts);
    auto out = torch::empty({N, 1, H, W}, opts);

    auto stream = at::cuda::getCurrentCUDAStream();

    group_stats_kernel<<<N * G, BLOCK_SIZE, 0, stream>>>(
        x_c.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
        C, G, HW, cpg, (float)eps);

    if (HW % 4 == 0) {
        int HW4 = HW / 4;
        long total4 = (long)N * HW4;
        int blocks = (int)((total4 + BLOCK_SIZE - 1) / BLOCK_SIZE);
        fused_norm_act_lse_vec4_kernel<<<blocks, BLOCK_SIZE, 0, stream>>>(
            reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
            mean.data_ptr<float>(), rstd.data_ptr<float>(),
            w_c.data_ptr<float>(), b_c.data_ptr<float>(),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            total4, C, G, HW4, cpg);
    } else {
        long total = (long)N * HW;
        int blocks = (int)((total + BLOCK_SIZE - 1) / BLOCK_SIZE);
        fused_norm_act_lse_kernel<<<blocks, BLOCK_SIZE, 0, stream>>>(
            x_c.data_ptr<float>(),
            mean.data_ptr<float>(), rstd.data_ptr<float>(),
            w_c.data_ptr<float>(), b_c.data_ptr<float>(),
            out.data_ptr<float>(),
            total, C, G, HW, cpg);
    }

    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_gn_act_lse(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, "
    "int64_t groups, double eps);"
)

fused_op = load_inline(
    name="fused_gn_act_lse_ext",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_source,
    functions=["fused_gn_act_lse"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: keeps the convolution in cuDNN, and replaces
    GroupNorm + Tanh + HardSwish + Residual Add + LogSumExp with custom
    fused CUDA kernels.
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, eps=1e-5):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.group_norm = nn.GroupNorm(groups, out_channels, eps=eps)
        self.groups = groups
        self.eps = eps

    def forward(self, x):
        # Convolution (cuDNN)
        x_conv = self.conv(x)
        # Fused GroupNorm + Tanh + HardSwish + Residual + LogSumExp
        return fused_op.fused_gn_act_lse(
            x_conv,
            self.group_norm.weight,
            self.group_norm.bias,
            self.groups,
            float(self.eps),
        )
```